Update app.py

app.py

@@ -1,11 +1,13 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
 import os
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# 初始化 Inference API 客戶端
+# 設置模型 ID 和加載 Hugging Face API token
 model_id = "dingckc/FineLlama-3.1-8B"
-api_token = os.getenv('ACCESS_KEY')
-inference = InferenceClient(model=model_id, token=api_token)
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=os.getenv('ACCESS_KEY'))
+model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=os.getenv('ACCESS_KEY'))
+model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
 # 定義推理函數
 def evaluate_essay(title, essay):
@@ -15,9 +17,10 @@ def evaluate_essay(title, essay):
     Essay: {essay}
     Please generate a detailed evaluation based on the rubric provided above.
     """
-    # 使用 text_generation 方法進行推理
-    response = inference.text_generation(input_text)
-    return response[0]["generated_text"] if "generated_text" in response[0] else "No evaluation available."
+    inputs = tokenizer(input_text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    with torch.no_grad():
+        outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=150)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # 使用 Gradio 構建界面
 title_input = gr.Textbox(label="Essay Title")
@@ -30,4 +33,4 @@ gr.Interface(
     outputs=output_text,
     title="Essay Evaluation",
     description="Enter the title and content of your essay to receive an evaluation."
-).launch()
+).launch(share=True)