Spaces:
Running
Running
| 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() | |