S-Dreamer commited on
Commit
e4233b9
·
verified ·
1 Parent(s): c57a0e5

Update app.py

Browse files

✅ What’s Improved
- Corrected missing import
- Robust error handling with feedback in UI
- Async model usage via asyncio.to_thread
- Modularized prompt building and response extraction
- Input validation
- Docstrings for every function
- Type annotations for clarity
- Clear comments and section separation
- Graceful fallback if model fails to load
- No blocking UI operations

Files changed (1) hide show
  1. app.py +93 -60
app.py CHANGED
@@ -1,17 +1,18 @@
1
- '''🔧 Prerequisites
 
 
 
 
 
2
 
3
- Install the necessary packages:
4
-
5
- pip install gradio transformers
6
-
7
- 📱 Gradio Chatbot App Code
8
- '''
9
  import gradio as gr
10
  from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
11
  import torch
 
 
12
 
13
- # List of available premium models
14
- premium_models = [
15
  "HuggingFaceH4/zephyr-7b-beta",
16
  "K00B404/BagOClownCoders-slerp-7B",
17
  "Qwen/Qwen2.5-Omni-7B",
@@ -21,57 +22,111 @@ premium_models = [
21
  "Alibaba-NLP/gte-Qwen2-7B-instruct",
22
  ]
23
 
24
- # Dictionary to cache loaded pipelines
25
- pipeline_cache = {}
 
 
 
26
 
27
- # Initial system prompt
28
- default_system_prompt = "You are a ChatBuddy and chat with the user in a Human way."
29
 
30
- def load_pipeline(model_name):
31
- if model_name not in pipeline_cache:
 
 
 
 
 
 
32
  print(f"Loading model: {model_name}")
33
  tokenizer = AutoTokenizer.from_pretrained(model_name)
34
- model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
35
- pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
 
 
 
 
 
 
 
36
  pipeline_cache[model_name] = pipe
37
- return pipeline_cache[model_name]
 
 
 
38
 
39
- def chatbot(user_input, history, model_choice):
40
- pipe = load_pipeline(model_choice)
41
 
42
- # Prepare the chat messages
43
- messages = [{"role": "system", "content": default_system_prompt}]
 
 
 
44
  for pair in history:
45
  messages.append({"role": "user", "content": pair[0]})
46
  messages.append({"role": "assistant", "content": pair[1]})
47
  messages.append({"role": "user", "content": user_input})
48
 
49
- # Flatten into a prompt string
50
  prompt = ""
51
  for msg in messages:
52
- if msg["role"] == "system":
53
- prompt += f"<|system|> {msg['content']}\n"
54
- elif msg["role"] == "user":
55
- prompt += f"<|user|> {msg['content']}\n"
56
- elif msg["role"] == "assistant":
57
- prompt += f"<|assistant|> {msg['content']}\n"
58
-
59
- # Generate a response
60
- response = pipe(prompt, max_new_tokens=200, do_sample=True, top_p=0.95, temperature=0.7)[0]['generated_text']
61
-
62
- # Extract only the last assistant response
63
- split_res = response.split("<|assistant|>")
64
- final_response = split_res[-1].strip() if len(split_res) > 1 else response
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
 
66
  history.append((user_input, final_response))
67
  return "", history
68
 
 
 
69
  with gr.Blocks() as demo:
70
  gr.Markdown("# 🤖 ChatBuddy - Advanced Chatbot with Selectable LLMs")
71
 
72
  with gr.Row():
73
- model_choice = gr.Dropdown(label="Select Model", choices=premium_models, value=premium_models[0])
74
-
 
 
75
  chatbot_ui = gr.Chatbot()
76
  user_input = gr.Textbox(show_label=False, placeholder="Type your message and press Enter")
77
  clear_btn = gr.Button("Clear")
@@ -82,25 +137,3 @@ with gr.Blocks() as demo:
82
  clear_btn.click(lambda: ([], ""), None, [chatbot_ui, state])
83
 
84
  demo.launch()
85
- '''
86
- ✅ Features:
87
-
88
- Model selection from dropdown
89
-
90
- Maintains chat history
91
-
92
- Respects a system prompt
93
-
94
- Uses text-generation pipeline
95
-
96
- 🧠 Optional Upgrades:
97
-
98
- Replace text-generation with chat-completion if models support it (like OpenChat, Mistral-instruct, etc.)
99
-
100
- Add streaming or token-by-token response if supported
101
-
102
- Save/load chat history
103
-
104
- Add support for vision models (Qwen2.5-VL-7B-Instruct) using a different UI tab
105
-
106
- '''
 
1
+ #!/usr/bin/env python3
2
+ # -*- coding: utf-8 -*-
3
+ """
4
+ Path: chatbot_app.py
5
+ Description: Gradio-based chatbot with selectable Hugging Face LLMs, using transformers pipelines.
6
+ """
7
 
 
 
 
 
 
 
8
  import gradio as gr
9
  from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
10
  import torch
11
+ import asyncio
12
+ from typing import List, Tuple, Dict
13
 
14
+ # List of available premium models
15
+ PREMIUM_MODELS = [
16
  "HuggingFaceH4/zephyr-7b-beta",
17
  "K00B404/BagOClownCoders-slerp-7B",
18
  "Qwen/Qwen2.5-Omni-7B",
 
22
  "Alibaba-NLP/gte-Qwen2-7B-instruct",
23
  ]
24
 
25
+ # Cache for loaded pipelines
26
+ pipeline_cache: Dict[str, pipeline] = {}
27
+
28
+ # ✅ Initial system prompt
29
+ DEFAULT_SYSTEM_PROMPT = "You are a ChatBuddy and chat with the user in a Human way."
30
 
 
 
31
 
32
+ def load_pipeline(model_name: str) -> pipeline:
33
+ """
34
+ Load and cache the text generation pipeline for the given model.
35
+ """
36
+ if model_name in pipeline_cache:
37
+ return pipeline_cache[model_name]
38
+
39
+ try:
40
  print(f"Loading model: {model_name}")
41
  tokenizer = AutoTokenizer.from_pretrained(model_name)
42
+ model = AutoModelForCausalLM.from_pretrained(
43
+ model_name,
44
+ torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
45
+ )
46
+ pipe = pipeline(
47
+ "text-generation",
48
+ model=model,
49
+ tokenizer=tokenizer,
50
+ device=0 if torch.cuda.is_available() else -1,
51
+ )
52
  pipeline_cache[model_name] = pipe
53
+ return pipe
54
+
55
+ except Exception as e:
56
+ raise RuntimeError(f"Failed to load model '{model_name}': {str(e)}")
57
 
 
 
58
 
59
+ def build_prompt(user_input: str, history: List[Tuple[str, str]]) -> str:
60
+ """
61
+ Construct the prompt string with system prompt, history, and current user input.
62
+ """
63
+ messages = [{"role": "system", "content": DEFAULT_SYSTEM_PROMPT}]
64
  for pair in history:
65
  messages.append({"role": "user", "content": pair[0]})
66
  messages.append({"role": "assistant", "content": pair[1]})
67
  messages.append({"role": "user", "content": user_input})
68
 
 
69
  prompt = ""
70
  for msg in messages:
71
+ role_tag = f"<|{msg['role']}|>"
72
+ prompt += f"{role_tag} {msg['content']}\n"
73
+
74
+ return prompt
75
+
76
+
77
+ def extract_response(generated_text: str) -> str:
78
+ """
79
+ Extract the last assistant response from generated text.
80
+ """
81
+ if "<|assistant|>" in generated_text:
82
+ split_res = generated_text.split("<|assistant|>")
83
+ return split_res[-1].strip()
84
+ return generated_text.strip()
85
+
86
+
87
+ async def chatbot(
88
+ user_input: str, history: List[Tuple[str, str]], model_choice: str
89
+ ) -> Tuple[str, List[Tuple[str, str]]]:
90
+ """
91
+ Main chatbot logic to generate model response asynchronously.
92
+ """
93
+ if not user_input.strip():
94
+ return "", history # Ignore empty inputs
95
+
96
+ try:
97
+ pipe = await asyncio.to_thread(load_pipeline, model_choice)
98
+ prompt = build_prompt(user_input, history)
99
+
100
+ response = await asyncio.to_thread(
101
+ pipe,
102
+ prompt,
103
+ max_new_tokens=200,
104
+ do_sample=True,
105
+ top_p=0.95,
106
+ temperature=0.7,
107
+ )
108
+
109
+ generated_text = response[0]["generated_text"]
110
+ final_response = extract_response(generated_text)
111
+
112
+ except RuntimeError as load_err:
113
+ final_response = str(load_err)
114
+ except Exception as e:
115
+ final_response = f"⚠️ Error during generation: {str(e)}"
116
 
117
  history.append((user_input, final_response))
118
  return "", history
119
 
120
+
121
+ # ✅ Gradio Interface
122
  with gr.Blocks() as demo:
123
  gr.Markdown("# 🤖 ChatBuddy - Advanced Chatbot with Selectable LLMs")
124
 
125
  with gr.Row():
126
+ model_choice = gr.Dropdown(
127
+ label="Select Model", choices=PREMIUM_MODELS, value=PREMIUM_MODELS[0]
128
+ )
129
+
130
  chatbot_ui = gr.Chatbot()
131
  user_input = gr.Textbox(show_label=False, placeholder="Type your message and press Enter")
132
  clear_btn = gr.Button("Clear")
 
137
  clear_btn.click(lambda: ([], ""), None, [chatbot_ui, state])
138
 
139
  demo.launch()