.

app.py

@@ -1,8 +1,20 @@
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def generate_reply(prompt):
-    return "您好！這是針對您輸入的回覆：" + prompt
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.float16, device_map="auto")
 
-demo = gr.Interface(fn=generate_reply, inputs="text", outputs="text")
+def respond(message):
+    prompt = f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=256)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # 擷取模型回覆部分
+    if "<start_of_turn>model" in response:
+        response = response.split("<start_of_turn>model")[-1].strip()
+    return response
 
-demo.launch()
+iface = gr.Interface(fn=respond, inputs="text", outputs="text")
+iface.launch()