Update app.py

app.py

@@ -1,44 +1,44 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 MODEL_ID = "caobin/llm-caobin"  
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True
+)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
-    device_map="auto"
+    device_map="auto",
+    trust_remote_code=True
 )
 
 def chat_fn(message, history):
-    input_text = ""
-
-    
+    full_prompt = ""
     for user_msg, bot_msg in history:
-        input_text += f"<|user|>{user_msg}<|assistant|>{bot_msg}"
-    input_text += f"<|user|>{message}<|assistant|>"
+        full_prompt += f"<|user|>{user_msg}<|assistant|>{bot_msg}"
+    full_prompt += f"<|user|>{message}<|assistant|>"
 
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
 
     output_ids = model.generate(
         **inputs,
         max_new_tokens=512,
-        do_sample=True,
         temperature=0.7,
         top_p=0.9,
+        do_sample=True,
     )
 
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-    # 只取 assistant 最新的回答
     if "<|assistant|>" in output_text:
         output_text = output_text.split("<|assistant|>")[-1]
 
-    return output_text
+    return output_text.strip()
 
 
-# Gradio UI
 with gr.Blocks(title="caobin LLM chatbot") as demo:
     gr.Markdown("# 🤖 caobin 自定义 LLM 对话 Demo")