import spaces import os import threading from collections import deque import plotly.graph_objs as go import pynvml import gradio as gr from huggingface_hub import snapshot_download from vptq.app_utils import get_chat_loop_generator models = [ { "name": "VPTQ-community/Meta-Llama-3.1-8B-Instruct-v8-k65536-65536-woft", "bits": "4 bits" }, { "name": "VPTQ-community/Meta-Llama-3.1-8B-Instruct-v8-k65536-256-woft", "bits": "3 bits" }, ] # Queues for storing historical data (saving the last 100 GPU utilization and memory usage values) gpu_util_history = deque(maxlen=100) mem_usage_history = deque(maxlen=100) def initialize_nvml(): """ Initialize NVML (NVIDIA Management Library). """ pynvml.nvmlInit() def get_gpu_info(): """ Get GPU utilization and memory usage information. Returns: dict: A dictionary containing GPU utilization and memory usage information. """ handle = pynvml.nvmlDeviceGetHandleByIndex(0) # Assuming a single GPU setup utilization = pynvml.nvmlDeviceGetUtilizationRates(handle) memory = pynvml.nvmlDeviceGetMemoryInfo(handle) gpu_info = { 'gpu_util': utilization.gpu, 'mem_used': memory.used / 1024**2, # Convert bytes to MiB 'mem_total': memory.total / 1024**2, # Convert bytes to MiB 'mem_percent': (memory.used / memory.total) * 100 } return gpu_info def update_charts(chart_height: int = 200) -> go.Figure: """ Update the GPU utilization and memory usage charts. Args: chart_height (int, optional): used to set the height of the chart. Defaults to 200. Returns: plotly.graph_objs.Figure: The updated figure containing the GPU and memory usage charts. """ # obtain GPU information gpu_info = get_gpu_info() # records the latest GPU utilization and memory usage values gpu_util = round(gpu_info.get('gpu_util', 0), 1) mem_used = round(gpu_info.get('mem_used', 0) / 1024, 2) # Convert MiB to GiB gpu_util_history.append(gpu_util) mem_usage_history.append(mem_used) # create GPU utilization line chart gpu_trace = go.Scatter( y=list(gpu_util_history), mode='lines+markers', text=list(gpu_util_history), line=dict(shape='spline', color='blue'), # Make the line smooth and set color yaxis='y1' # Link to y-axis 1 ) # create memory usage line chart mem_trace = go.Scatter( y=list(mem_usage_history), mode='lines+markers', text=list(mem_usage_history), line=dict(shape='spline', color='red'), # Make the line smooth and set color yaxis='y2' # Link to y-axis 2 ) # set the layout of the chart layout = go.Layout( xaxis=dict(title=None, showticklabels=False, ticks=''), yaxis=dict( title='GPU Utilization (%)', range=[-5, 110], titlefont=dict(color='blue'), tickfont=dict(color='blue'), ), yaxis2=dict(title='Memory Usage (GiB)', range=[0, max(24, max(mem_usage_history) + 1)], titlefont=dict(color='red'), tickfont=dict(color='red'), overlaying='y', side='right'), height=chart_height, # set the height of the chart margin=dict(l=10, r=10, t=0, b=0), # set the margin of the chart showlegend=False # disable the legend ) fig = go.Figure(data=[gpu_trace, mem_trace], layout=layout) return fig def initialize_history(): """ Initializes the GPU utilization and memory usage history. """ for _ in range(100): gpu_info = get_gpu_info() gpu_util_history.append(round(gpu_info.get('gpu_util', 0), 1)) mem_usage_history.append(round(gpu_info.get('mem_percent', 0), 1)) def enable_gpu_info(): pynvml.nvmlInit() def disable_gpu_info(): pynvml.nvmlShutdown() model_choices = [f"{model['name']} ({model['bits']})" for model in models] display_to_model = {f"{model['name']} ({model['bits']})": model['name'] for model in models} def download_model(model): print(f"Downloading {model['name']}...") snapshot_download(repo_id=model['name']) def download_models_in_background(): print('Downloading models for the first time...') for model in models: download_model(model) download_thread = threading.Thread(target=download_models_in_background) download_thread.start() loaded_models = {} @spaces.GPU def respond( message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p, selected_model_display_label, ): model_name = display_to_model[selected_model_display_label] # Check if the model is already loaded if model_name not in loaded_models: # Load and store the model in the cache loaded_models[model_name] = get_chat_loop_generator(model_name) chat_completion = loaded_models[model_name] messages = [{"role": "system", "content": system_message}] for val in history: if val[0]: messages.append({"role": "user", "content": val[0]}) if val[1]: messages.append({"role": "assistant", "content": val[1]}) messages.append({"role": "user", "content": message}) response = "" for message in chat_completion( messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p, ): token = message response += token yield response """ For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface """ # enable_gpu_info() with gr.Blocks(fill_height=True) as demo: with gr.Row(): def update_chart(): return _update_charts(chart_height=200) gpu_chart = gr.Plot(update_chart, every=0.1) # update every 0.1 seconds with gr.Column(): chat_interface = gr.ChatInterface( respond, additional_inputs=[ gr.Textbox(value="You are a friendly Chatbot.", label="System message"), gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"), gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), gr.Slider( minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)", ), gr.Dropdown( choices=model_choices, value=model_choices[0], label="Select Model", ), ], ) if __name__ == "__main__": share = os.getenv("SHARE_LINK", None) in ["1", "true", "True"] demo.launch(share=share) # disable_gpu_info()