Update app.py

app.py

@@ -1,13 +1,6 @@
 import gradio as gr
 import os
-# api_token = os.getenv("HF_TOKEN")
-
-from dotenv import load_dotenv
-
-from langchain.llms.base import LLM
-from transformers import AutoTokenizer
-from huggingface_hub import HfApi
-import requests
+api_token = os.getenv("HF_TOKEN")
 
 
 from langchain_community.vectorstores import FAISS
@@ -19,67 +12,12 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
-from langchain_community.llms import HuggingFaceHub, HuggingFaceEndpoint
+from langchain_community.llms import HuggingFaceEndpoint
 import torch
 
-
-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 = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]  
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-def retrieve_api():
-    """Retrieve HuggingFace API Key"""
-    _ = load_dotenv()
-    global huggingfacehub_api_token
-    api_token = os.environ.get("HUGGINGFACE_API_KEY")
-
-# 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
@@ -105,34 +43,18 @@ 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(
-    #         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":
+    if llm_model == "meta-llama/Meta-Llama-3-8B-Instruct":
         llm = HuggingFaceEndpoint(
             repo_id=llm_model,
-            task="text-generation",
             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",
-        #     huggingfacehub_api_token=api_token,
-        #     model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens}
-        #     )
-        
     else:
         llm = HuggingFaceEndpoint(
             huggingfacehub_api_token = api_token,
-            repo_id=llm_model,
-            task="text-generation",
+            repo_id=llm_model, 
             temperature = temperature,
             max_new_tokens = max_tokens,
             top_k = top_k,
@@ -291,5 +213,4 @@ def demo():
 
 
 if __name__ == "__main__":
-    retrieve_api()
-    demo()
+    demo()