v6

app.py

@@ -1,20 +1,24 @@
-import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from huggingface_hub import login
+import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+# 登入 Hugging Face（換成你自己的 token）
+hf_token = os.environ.get("HF_TOKEN")
+login(hf_token)
 
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.float16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-2b-it",
+    torch_dtype=torch.float32
+).to("cpu")
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-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
+def generate(text):
+    result = pipe(text, max_new_tokens=200, do_sample=True, temperature=0.7)
+    return result[0]["generated_text"]
 
-iface = gr.Interface(fn=respond, inputs="text", outputs="text")
-iface.launch()
+demo = gr.Interface(fn=generate, inputs="text", outputs="text")
+demo.launch()