Update app.py

app.py

@@ -2,6 +2,11 @@ import gradio as gr
 import os
 api_token = os.getenv("HF_TOKEN")
 
+from langchain.llms.base import LLM
+from transformers import AutoTokenizer
+from huggingface_hub import HfApi
+import requests
+
 
 from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import PyPDFLoader
@@ -15,9 +20,64 @@ from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceHub, HuggingFaceEndpoint
 import torch
 
-list_llm = ["meta-llama/Llama-3.1-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]   # meta-llama/Meta-Llama-3-8B-Instruct
+
+from langchain.llms.base import LLM
+from transformers import AutoTokenizer
+from huggingface_hub import HfApi
+import requests
+
+list_llm = ["HuggingFaceH4/zephyr-7b-beta", "meta-llama/Llama-3.1-8B-Instruct"] # "mistralai/Mistral-7B-Instruct-v0.2"   # meta-llama/Meta-Llama-3-8B-Instruct
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
+class ZephyrLLM(LLM):
+    def __init__(self, repo_id, huggingfacehub_api_token, max_new_tokens=512, temperature=0.7, **kwargs):
+        super().__init__(**kwargs)
+        self.repo_id = repo_id
+        self.api_token = huggingfacehub_api_token
+        self.api_url = f"https://api-inference.huggingface.co/models/{repo_id}"
+        self.headers = {"Authorization": f"Bearer {huggingfacehub_api_token}"}
+        self.tokenizer = AutoTokenizer.from_pretrained(repo_id)
+        self.max_new_tokens = max_new_tokens
+        self.temperature = temperature
+
+    def _call(self, prompt, stop=None):
+    # Format as chat message
+    messages = [{"role": "user", "content": prompt}]
+    
+    # Apply Zephyr's chat template
+    formatted_prompt = self.tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    
+    # Send request to Hugging Face Inference API
+    payload = {
+        "inputs": formatted_prompt,
+        "parameters": {
+            "max_new_tokens": self.max_new_tokens,
+            "temperature": self.temperature
+        }
+    }
+    response = requests.post(self.api_url, headers=self.headers, json=payload)
+
+    if response.status_code == 200:
+        full_response = response.json()[0]["generated_text"]
+
+        # Extract the assistant reply from the full response
+        # After <|assistant|>\n, everything is the model's answer
+        if "<|assistant|>" in full_response:
+            return full_response.split("<|assistant|>")[-1].strip()
+        else:
+            return full_response.strip()
+
+    else:
+        raise Exception(f"Failed call [{response.status_code}]: {response.text}")
+
+
+    @property
+    def _llm_type(self) -> str:
+        return "zephyr-custom"
+
+
 # Load and split PDF document
 def load_doc(list_file_path):
     # Processing for one document only
@@ -43,7 +103,14 @@ def create_db(splits):
 
 # Initialize langchain LLM chain
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    if llm_model == "meta-llama/Llama-3.1-8B-Instruct":
+    if llm_model == "HuggingFaceH4/zephyr-7b-beta":
+        llm = ZephyrLLM(
+            repo_id=llm_model,
+            huggingfacehub_api_token=api_token,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+        )
+    # if llm_model == "meta-llama/Llama-3.1-8B-Instruct":
         # llm = HuggingFaceEndpoint(
         #     repo_id=llm_model,
         #     huggingfacehub_api_token = api_token,
@@ -52,11 +119,11 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         #     top_k = top_k,
         # )
 
-        llm = HuggingFaceHub(
-            repo_id="mistralai/Mistral-7B-Instruct-v0.2",
-            huggingfacehub_api_token=api_token,
-            model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens}
-            )
+        # llm = HuggingFaceHub(
+        #     repo_id="mistralai/Mistral-7B-Instruct-v0.2",
+        #     huggingfacehub_api_token=api_token,
+        #     model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens}
+        #     )
         
     else:
         llm = HuggingFaceEndpoint(