Spaces:
Running
Running
File size: 2,849 Bytes
1894678 e467758 e47d2e4 e467758 1894678 505a5b8 1894678 e47d2e4 0be02c9 1894678 505a5b8 1894678 |
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 |
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
# List of available premium models
premium_models = [
"Qwen/Qwen2-1.5B-Instruct",
]
# Dictionary to cache loaded pipelines
pipeline_cache = {}
# Initial system prompt
default_system_prompt = "You are a ChatBuddy and chat with the user in a Human way."
def load_pipeline(model_name):
if model_name not in pipeline_cache:
print(f"Loading model: {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
pipeline_cache[model_name] = pipe
return pipeline_cache[model_name]
def chatbot(user_input, history, model_choice):
pipe = load_pipeline(model_choice)
# Prepare the chat messages
messages = [{"role": "system", "content": default_system_prompt}]
for pair in history:
messages.append({"role": "user", "content": pair[0]})
messages.append({"role": "assistant", "content": pair[1]})
messages.append({"role": "user", "content": user_input})
# Flatten into a prompt string
prompt = ""
for msg in messages:
if msg["role"] == "system":
prompt += f"<|system|> {msg['content']}\n"
elif msg["role"] == "user":
prompt += f"<|user|> {msg['content']}\n"
elif msg["role"] == "assistant":
prompt += f"<|assistant|> {msg['content']}\n"
# Generate a response
response = pipe(prompt, max_new_tokens=200, do_sample=True, top_p=0.95, temperature=0.7)[0]['generated_text']
# Extract only the last assistant response
split_res = response.split("<|assistant|>")
final_response = split_res[-1].strip() if len(split_res) > 1 else response
history.append({"role": "user", "content": user_input})
history.append({"role": "assistant", "content": final_response})
return "", history
with gr.Blocks() as demo:
gr.Markdown("# 🤖 ChatBuddy - Advanced Chatbot with Selectable LLMs")
with gr.Row():
model_choice = gr.Dropdown(label="Select Model", choices=premium_models, value=premium_models[0])
with gr.Row():
model_choice = gr.Textbox(label="System Prompt", value=default_system_prompt)
chatbot_ui = gr.Chatbot(type="messages")
user_input = gr.Textbox(show_label=False, placeholder="Type your message and press Enter")
clear_btn = gr.Button("Clear")
state = gr.State([])
user_input.submit(chatbot, [user_input, state, model_choice], [user_input, chatbot_ui])
clear_btn.click(lambda: ([], ""), None, [chatbot_ui, state])
demo.launch()
|