Update app.py

app.py

@@ -1,68 +1,102 @@
 import gradio as gr
-from gradio import Interface, Textbox, Image, Markdown
-from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPProcessor, CLIPModel
 import torch
-from PIL import Image as PILImage
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PyPDF2 import PdfReader
+import google.generativeai as genai
+import os
+from langsmith import Client
+from ragas.metrics import (
+    faithfulness, answer_relevancy, context_relevancy
+)
 
-# Load models and tokenizers
-openelm_270m_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-hf")
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+# 加載模型
+openelm_model = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M", trust_remote_code=True)
+openelm_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-hf")
 
-def generate_text(prompt, num_tokens):
-    tokenized_prompt = tokenizer(prompt, return_tensors="pt")
-    
-    output_ids = openelm_270m_instruct.generate(
+# Gemini API 設置
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+genai.configure(api_key=GOOGLE_API_KEY)
+
+# LangSmith 設置
+os.environ["LANGCHAIN_API_KEY"] = "your_langchain_api_key"
+os.environ["LANGCHAIN_TRACING_V2"] = "true"
+os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
+client = Client()
+
+def extract_text_from_pdf(pdf_path):
+    with open(pdf_path, 'rb') as file:
+        reader = PdfReader(file)
+        text = ""
+        for page in reader.pages:
+            text += page.extract_text() + "\n"
+    return text
+
+def gemini_generate(prompt, max_tokens):
+    model = genai.GenerativeModel('gemini-pro')
+    response = model.generate_content(prompt, max_output_tokens=max_tokens)
+    return response.text
+
+def nvidia_generate(prompt, max_tokens):
+    # 這裡需要實現 Nvidia API 調用
+    return "Nvidia API 尚未實現"
+
+def openelm_generate(prompt, max_tokens):
+    tokenized_prompt = openelm_tokenizer(prompt, return_tensors="pt")
+    output_ids = openelm_model.generate(
         tokenized_prompt["input_ids"],
-        max_length=int(num_tokens),
+        max_length=max_tokens,
         pad_token_id=0,
     )
-    
-    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return output_text
+    return openelm_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+def evaluate_response(response, context, query):
+    # 使用 RAGAS 評估回答
+    faith_score = faithfulness.score([response], [context], [query])
+    ans_rel_score = answer_relevancy.score([response], [query])
+    ctx_rel_score = context_relevancy.score([response], [context], [query])
+    return faith_score, ans_rel_score, ctx_rel_score
 
-def process_image(image):
-    if isinstance(image, str):
-        image = PILImage.open(image).convert("RGB")
-    inputs = clip_processor(images=image, return_tensors="pt")
-    image_features = clip_model.get_image_features(**inputs)
-    return image_features
+def process_query(pdf_file, llm_choice, query, max_tokens, api_key):
+    global GOOGLE_API_KEY
+    if api_key:
+        GOOGLE_API_KEY = api_key
+        genai.configure(api_key=GOOGLE_API_KEY)
 
-def generate_multimodal(text_prompt, image, num_tokens):
-    # Process the image
-    image_features = process_image(image)
+    # 從 PDF 提取文本
+    pdf_path = pdf_file.name
+    context = extract_text_from_pdf(pdf_path)
     
-    # Combine text prompt with image features (you may need to adjust this based on OpenELM's specific requirements)
-    combined_input = f"{text_prompt} [IMAGE]"
+    # 根據選擇的 LLM 生成回應
+    if llm_choice == "Gemini":
+        response = gemini_generate(f"上下文: {context}\n問題: {query}", max_tokens)
+    elif llm_choice == "Nvidia":
+        response = nvidia_generate(f"上下文: {context}\n問題: {query}", max_tokens)
+    else:  # OpenELM
+        response = openelm_generate(f"上下文: {context}\n問題: {query}", max_tokens)
     
-    # Generate text based on the combined input
-    output = generate_text(combined_input, num_tokens)
+    # 評估回應
+    faith_score, ans_rel_score, ctx_rel_score = evaluate_response(response, context, query)
     
-    return output
-
-def greet(text_input, image_input, num_tokens):
-    if image_input is not None:
-        return generate_multimodal(text_input, image_input, num_tokens)
-    else:
-        return generate_text(text_input, num_tokens)
-
-developer_info = """
-This space is developed by Ahmadreza Anaami
-Feel free to set via API key too
-Models used: apple/OpenELM-270M, openai/clip-vit-base-patch32
-"""
+    return response, faith_score, ans_rel_score, ctx_rel_score
 
+# Gradio 介面
 iface = gr.Interface(
-    fn=greet,
+    fn=process_query,
     inputs=[
-        Textbox(label="Enter Text Here:", type="text"),
-        Image(label="Upload Image (optional):"),
-        Textbox(label="Number of generated tokens:", type="text")
+        gr.File(label="上傳 PDF"),
+        gr.Dropdown(["Gemini", "Nvidia", "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="上下文相關性得分")
     ],
-    outputs=[Textbox(label="Generated answer:")],
-    title="OpenELM-270M Multimodal",
-    description=developer_info,
+    title="多模型 LLM 查詢介面，支持 PDF 上下文",
+    description="上傳 PDF，選擇 LLM，並提出問題。回應將使用 RAGAS 指標進行評估。",
     css="""
         #dev-info {
             font-size: 0.8rem;