Update app.py

app.py

@@ -4,12 +4,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
 
 MODEL_ID = "LMSeed/GPT2-small-distilled-100M"
-# HF_TOKEN = os.environ.get("HF_TOKEN")  # 如果模型私有
 
 device = 0 if torch.cuda.is_available() else -1
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+
 if torch.cuda.is_available():
     model = model.to("cuda")
 
@@ -17,30 +17,49 @@ generator = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=0 if torch.cuda.is_available() else -1
+    device=device
 )
 
 def chat_with_model(user_message, chat_history, max_new_tokens=60, temperature=0.8, top_p=0.9):
+
+    if chat_history is None:
+        chat_history = []
+
+    # Build conversation context
     history_text = ""
-    if chat_history:
-        for i, msg in enumerate(chat_history):
-            role = "User" if i % 2 == 0 else "Assistant"
-            history_text += f"{role}: {msg}\n"
+    for i, msg in enumerate(chat_history):
+        role = "User" if i % 2 == 0 else "Assistant"
+        history_text += f"{role}: {msg}\n"
+
     history_text += f"User: {user_message}\nAssistant: "
 
-    outputs = generator(history_text, max_new_tokens=max_new_tokens,
-                        do_sample=True, temperature=float(temperature),
-                        top_p=float(top_p), num_return_sequences=1)
+    # Generate
+    outputs = generator(
+        history_text,
+        max_new_tokens=int(max_new_tokens),
+        do_sample=True,
+        temperature=float(temperature),
+        top_p=float(top_p),
+        num_return_sequences=1
+    )
+
     reply = outputs[0]["generated_text"][len(history_text):].strip()
+
+    # Prevent model from continuing system formatting
     if "\n" in reply:
         reply = reply.split("\n")[0].strip()
-    chat_history = chat_history or []
+
+    # Update chat history
     chat_history.append(user_message)
     chat_history.append(reply)
-    return "", chat_history
+
+    return "", chat_history, chat_history
+
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Chat with GPT-2")
+
+    gr.Markdown("# Chat with Stu")
+
     with gr.Row():
         with gr.Column(scale=3):
             chat = gr.Chatbot(elem_id="chatbot", label="Conversation")
@@ -49,10 +68,16 @@ with gr.Blocks() as demo:
             max_tokens = gr.Slider(10, 256, value=60, label="max_new_tokens")
             temp = gr.Slider(0.1, 1.2, value=0.8, label="temperature")
             top_p = gr.Slider(0.1, 1.0, value=0.9, label="top_p")
+
         with gr.Column(scale=1):
             gr.Markdown("Model: " + MODEL_ID)
 
     state = gr.State([])
-    send.click(fn=chat_with_model, inputs=[msg, state, max_tokens, temp, top_p],
-               outputs=[msg, state])
+
+    send.click(
+        fn=chat_with_model,
+        inputs=[msg, state, max_tokens, temp, top_p],
+        outputs=[msg, chat, state]
+    )
+
     demo.launch()