Update app.py

app.py

@@ -3,25 +3,37 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_name = "winninghealth/WiNGPT-Babel"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True) #, torch_dtype="auto", device_map="auto"
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
 
 def translate(text):
     prompt = f"<|im_start|>system\n中英互译下面的内容<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
-    inputs = tokenizer([prompt], return_tensors="pt") #.to(model.device)
+    inputs = tokenizer([prompt], return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
     translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return translated_text
 
-def predict(text):
-    return translate(text)
+def custom_api(text_list, source_lang, target_lang):
+    # 假设你的模型只支持中英互译
+    if source_lang == "zh-CN" and target_lang == "en":
+        translated_list = [translate(text) for text in text_list]
+    elif source_lang == "en" and target_lang == "zh-CN":
+        translated_list = [translate(text) for text in text_list]
+    else:
+        return {"error": "Unsupported language pair"}
 
+    return {"translations": [{"detected_source_lang": source_lang, "text": translated_text} for translated_text in translated_list]}
+
+# 创建 Gradio 接口
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(lines=5, label="输入文本 (支持中英互译)"),
-    outputs=gr.Textbox(label="翻译结果"),
+    fn=custom_api,
+    inputs=[
+        gr.Textbox(lines=5, label="输入文本列表 (支持中英互译)", placeholder='["Hello", "World"]'),
+        gr.Textbox(label="源语言", placeholder="zh-CN"),
+        gr.Textbox(label="目标语言", placeholder="en")
+    ],
+    outputs=gr.JSON(label="翻译结果"),
     title="WiNGPT-Babel 翻译 Demo",
     description="基于 WiNGPT-Babel 模型的翻译演示。支持中英互译。",
-    examples=[["Hello, world!"], ["你好，世界！"]],
 )
 
 iface.launch()