Update backend/query_llm.py

backend/query_llm.py

@@ -9,8 +9,9 @@ from typing import Any, Dict, Generator, List
 from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer
 
-#tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
+#tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
+#tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x22B-Instruct-v0.1")
 temperature = 0.5
 top_p = 0.7
 repetition_penalty = 1.2
@@ -18,16 +19,16 @@ repetition_penalty = 1.2
 OPENAI_KEY = getenv("OPENAI_API_KEY")
 HF_TOKEN = getenv("HUGGING_FACE_HUB_TOKEN")
 
-#hf_client = InferenceClient(
-#        "mistralai/Mistral-7B-Instruct-v0.1",
-#        token=HF_TOKEN
-#        )
+hf_client = InferenceClient(
+       "mistralai/Mistral-7B-Instruct-v0.1",
+       token=HF_TOKEN
+       )
 
 
-hf_client = InferenceClient(
-        "mistralai/Mixtral-8x7B-Instruct-v0.1",
-        token=HF_TOKEN
-        )
+# hf_client = InferenceClient(
+#         "mistralai/Mixtral-8x7B-Instruct-v0.1",
+#         token=HF_TOKEN
+#         )
 def format_prompt(message: str, api_kind: str):
     """
     Formats the given message using a chat template.