File size: 3,545 Bytes
d0f4aff 5968a97 2db3bb3 d0f4aff 5968a97 5e84c69 4b15ccd 2db3bb3 b6b7c74 0ebe852 b6b7c74 2db3bb3 b6b7c74 65d5ebe b6b7c74 5968a97 2db3bb3 5968a97 2db3bb3 5968a97 4b15ccd 5968a97 2db3bb3 5968a97 2db3bb3 b6b7c74 5e84c69 b16f2d9 0115682 d0f4aff 2db3bb3 4b15ccd 2db3bb3 5e84c69 2db3bb3 b6b7c74 2db3bb3 0115682 2db3bb3 0115682 2db3bb3 4b15ccd 0ebe852 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import os
import torch
import time
import gradio as gr
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM
# Globals
current_model = None
current_tokenizer = None
def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
global current_model, current_tokenizer
token = os.getenv("HF_TOKEN")
progress(0, desc="Loading tokenizer...")
current_tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
progress(0.5, desc="Loading model...")
current_model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="cpu", # loaded to CPU initially
use_auth_token=token
)
progress(1, desc="Model ready.")
return f"{model_name} loaded and ready!"
# Format conversation as plain text
def format_prompt(messages):
prompt = ""
for msg in messages:
role = msg["role"]
if role == "user":
prompt += f"User: {msg['content'].strip()}\n"
elif role == "assistant":
prompt += f"Assistant: {msg['content'].strip()}\n"
prompt += "Assistant:"
return prompt
def add_user_message(user_input, history):
return "", history + [{"role": "user", "content": user_input}]
# Available models
model_choices = [
"meta-llama/Llama-3.2-3B-Instruct",
"deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
"google/gemma-7b"
]
# UI
with gr.Blocks() as demo:
gr.Markdown("## Clinical Chatbot (Streaming) — LLaMA, DeepSeek, Gemma")
default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
@spaces.GPU
def chat_with_model(messages):
global current_model, current_tokenizer
if current_model is None or current_tokenizer is None:
yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
return
current_model.to("cuda")
prompt = format_prompt(messages)
inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
output_ids = []
messages = messages.copy()
messages.append({"role": "assistant", "content": ""})
for token_id in current_model.generate(
**inputs,
max_new_tokens=256,
do_sample=False,
return_dict_in_generate=True,
output_scores=False
).sequences[0][inputs['input_ids'].shape[-1]:]: # skip input tokens
output_ids.append(token_id.item())
decoded = current_tokenizer.decode(output_ids, skip_special_tokens=True)
messages[-1]["content"] = decoded
yield messages
with gr.Row():
model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
model_status = gr.Textbox(label="Model Status", interactive=False)
chatbot = gr.Chatbot(label="Chat", type="messages")
msg = gr.Textbox(label="Your message", placeholder="Enter clinical input...", show_label=False)
clear = gr.Button("Clear")
# Load default model on startup
demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)
# Load selected model manually
model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)
# Submit message + stream model response
msg.submit(add_user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
chat_with_model, chatbot, chatbot
)
# Clear chat
clear.click(lambda: [], None, chatbot, queue=False)
demo.launch()
|