add different llms

app.py

@@ -63,25 +63,25 @@ def bot(history, api_kind, chunk_table, embedding_model, llm_model, cross_encode
     prompt_html = template_html.render(documents=documents, query=query)
 
     if llm_model == "mistralai/Mistral-7B-Instruct-v0.2":
-        pass
+        generate_fn = generate_hf
     if llm_model == "mistralai/Mistral-7B-v0.1":
-        pass
+        generate_fn = generate_hf
     if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
-        pass 
+        generate_fn = generate_hf 
     if llm_model == "gpt-3.5-turbo":
-        pass
+        generate_fn = generate_openai
     if llm_model == "gpt-4-turbo-preview":
-        pass
+        generate_fn = generate_openai
 
     #if api_kind == "HuggingFace":
     #     generate_fn = generate_hf
     #elif api_kind == "OpenAI":
     #     generate_fn = generate_openai
     #else:
-    #     raise gr.Error(f"API {api_kind} is not supported")
+         raise gr.Error(f"API {api_kind} is not supported")
 
     history[-1][1] = ""
-    for character in generate_fn(prompt, history[:-1]):
+    for character in generate_fn(prompt, history[:-1], llm_model):
         history[-1][1] = character
         yield history, prompt_html
 

backend/query_llm.py

@@ -34,7 +34,7 @@ OAI_GENERATE_KWARGS = {
 }
 
 
-def format_prompt(message: str, api_kind: str):
+def format_prompt(message: str, api_kind: str, tokenizer_name = None):
     """
     Formats the given message using a chat template.
 
@@ -51,12 +51,13 @@ def format_prompt(message: str, api_kind: str):
     if api_kind == "openai":
         return messages
     elif api_kind == "hf":
+        TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name)
         return TOKENIZER.apply_chat_template(messages, tokenize=False)
     elif api_kind:
         raise ValueError("API is not supported")
 
 
-def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
+def generate_hf(prompt: str, history: str, hf_model_name: str) -> Generator[str, None, str]:
     """
     Generate a sequence of tokens based on a given prompt and history using Mistral client.
 
@@ -67,8 +68,14 @@ def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
         Generator[str, None, str]: A generator yielding chunks of generated text.
                                    Returns a final string if an error occurs.
     """
+    
 
-    formatted_prompt = format_prompt(prompt, "hf")
+    HF_CLIENT = InferenceClient(
+        hf_model_name,
+        token=os.getenv("HF_TOKEN")
+    )
+
+    formatted_prompt = format_prompt(prompt, "hf", hf_model_name)
     formatted_prompt = formatted_prompt.encode("utf-8").decode("utf-8")
 
     try:
@@ -93,7 +100,7 @@ def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
             raise gr.Error(f"Unhandled Exception: {str(e)}")
 
 
-def generate_openai(prompt: str, history: str) -> Generator[str, None, str]:
+def generate_openai(prompt: str, history: str, model_name: str) -> Generator[str, None, str]:
     """
     Generate a sequence of tokens based on a given prompt and history using Mistral client.
 
@@ -108,7 +115,8 @@ def generate_openai(prompt: str, history: str) -> Generator[str, None, str]:
 
     try:
         stream = OAI_CLIENT.chat.completions.create(
-            model=os.getenv("OPENAI_MODEL"),
+            #model=os.getenv("OPENAI_MODEL"),
+            model = model_name,
             messages=formatted_prompt,
             **OAI_GENERATE_KWARGS,
             stream=True