Update app.py

app.py

@@ -26,55 +26,55 @@ 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 = ["meta-llama/Llama-3.1-8B-Instruct"] # , "HuggingFaceH4/zephyr-7b-beta"] # "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()
+# 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}")
+    #     else:
+    #         raise Exception(f"Failed call [{response.status_code}]: {response.text}")
 
 
-    @property
-    def _llm_type(self) -> str:
-        return "zephyr-custom"
+    # @property
+    # def _llm_type(self) -> str:
+    #     return "zephyr-custom"
 
 
 # Load and split PDF document
@@ -102,21 +102,22 @@ 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 == "HuggingFaceH4/zephyr-7b-beta":
-        llm = ZephyrLLM(
+    # 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,
-            temperature=temperature,
-            max_new_tokens=max_tokens,
+            task="text-generation",
+            huggingfacehub_api_token = api_token,
+            temperature = temperature,
+            max_new_tokens = max_tokens,
+            top_k = top_k,
         )
-    # if llm_model == "meta-llama/Llama-3.1-8B-Instruct":
-        # llm = HuggingFaceEndpoint(
-        #     repo_id=llm_model,
-        #     huggingfacehub_api_token = api_token,
-        #     temperature = temperature,
-        #     max_new_tokens = max_tokens,
-        #     top_k = top_k,
-        # )
 
         # llm = HuggingFaceHub(
         #     repo_id="mistralai/Mistral-7B-Instruct-v0.2",
@@ -127,7 +128,8 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     else:
         llm = HuggingFaceEndpoint(
             huggingfacehub_api_token = api_token,
-            repo_id=llm_model, 
+            repo_id=llm_model,
+            task="text-generation",
             temperature = temperature,
             max_new_tokens = max_tokens,
             top_k = top_k,