Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, GPT2Tokenizer
 from PyPDF2 import PdfReader
 import google.generativeai as genai
 import os
@@ -8,8 +8,8 @@ from langsmith import Client
 from ragas.metrics import faithfulness, answer_relevancy, context_relevancy
 
 # 加載模型
-openelm_model = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M", trust_remote_code=True)
-openelm_tokenizer = AutoTokenizer.from_pretrained("apple/OpenELM-270M", trust_remote_code=True)  # Add trust_remote_code here
+openelm_model = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M", revision="main", trust_remote_code=True)
+openelm_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")  # OpenELM 使用 GPT2 tokenizer
 
 # Gemini API 設置
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
@@ -35,11 +35,11 @@ def gemini_generate(prompt, max_tokens):
     return response.text
 
 def openelm_generate(prompt, max_tokens):
-    tokenized_prompt = openelm_tokenizer(prompt, return_tensors="pt")
+    input_ids = openelm_tokenizer.encode(prompt, return_tensors="pt")
     output_ids = openelm_model.generate(
-        tokenized_prompt["input_ids"],
+        input_ids,
         max_length=max_tokens,
-        pad_token_id=0,
+        pad_token_id=openelm_tokenizer.eos_token_id
     )
     return openelm_tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
@@ -56,42 +56,21 @@ def process_query(pdf_file, llm_choice, query, max_tokens, api_key):
             GOOGLE_API_KEY = api_key
             genai.configure(api_key=GOOGLE_API_KEY)
         
-        # 從 PDF 提取文本
         pdf_path = pdf_file.name
         context = extract_text_from_pdf(pdf_path)
         
-        # 根據選擇的 LLM 生成回應
         if llm_choice == "Gemini":
             response = gemini_generate(f"上下文: {context}\n問題: {query}", max_tokens)
         else:  # OpenELM
             response = openelm_generate(f"上下文: {context}\n問題: {query}", max_tokens)
         
-        # 評估回應
         faith_score, ans_rel_score, ctx_rel_score = evaluate_response(response, context, query)
         
         return response, faith_score, ans_rel_score, ctx_rel_score
     except Exception as e:
-        return str(e), 0, 0, 0  # 返回錯誤消息和零分數
+        return str(e), 0, 0, 0
 
-# Gradio 介面
-iface = gr.Interface(
-    fn=process_query,
-    inputs=[
-        gr.File(label="上傳 PDF"),
-        gr.Dropdown(["Gemini", "OpenELM"], label="選擇 LLM"),
-        gr.Textbox(label="輸入您的問題"),
-        gr.Slider(minimum=50, maximum=1000, step=50, label="最大令牌數"),
-        gr.Textbox(label="Gemini API Key (可選)", type="password")
-    ],
-    outputs=[
-        gr.Textbox(label="生成的答案"),
-        gr.Number(label="真實性得分"),
-        gr.Number(label="答案相關性得分"),
-        gr.Number(label="上下文相關性得分")
-    ],
-    title="多模型 LLM 查詢介面，支持 PDF 上下文",
-    description="上傳 PDF，選擇 LLM，並提出問題。回應將使用 RAGAS 指標進行評估。"
-)
+# Gradio 界面設置保持不變...
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()