change complition model interface

app.py

@@ -39,6 +39,14 @@ def bot(history, chunk_table, embedding_model, llm_model, cross_encoder, top_k_p
     top_k_param = int(top_k_param)
     query = history[-1][0]
 
+    logger.info("bot launched ...")
+    logger.info(f"embedding model: {embedding_model}")
+    logger.info(f"LLM model: {llm_model}")
+    logger.info(f"Cross encoder model: {cross_encoder}")
+    logger.info(f"TopK: {top_k_param}")
+    logger.info(f"ReRank TopK: {rerank_topk}")
+
+
     if not query:
         raise gr.Warning("Please submit a non-empty string as a prompt")
 
@@ -48,9 +56,13 @@ def bot(history, chunk_table, embedding_model, llm_model, cross_encoder, top_k_p
 
     #documents = retrieve(query, TOP_K)
     documents = retrieve(query, top_k_param, chunk_table, embedding_model)
+    logger.info('Retrived document count:', len(documents))
+
     if cross_encoder != "None" and len(documents) > 1:
         documents = rerank_documents(cross_encoder, documents, query, top_k_rerank=rerank_topk)
         #"cross-encoder/ms-marco-MiniLM-L-6-v2"
+        logger.info('ReRank done, document count:', len(documents))
+
 
 
 
@@ -79,7 +91,8 @@ def bot(history, chunk_table, embedding_model, llm_model, cross_encoder, top_k_p
     #     generate_fn = generate_openai
     #else:
     #     raise gr.Error(f"API {api_kind} is not supported")
-
+    
+    logger.info(f'Complition started. llm_model: {llm_model}, prompt: {prompt}')
     history[-1][1] = ""
     for character in generate_fn(prompt, history[:-1], llm_model):
         history[-1][1] = character

backend/query_llm.py

@@ -10,12 +10,12 @@ from transformers import AutoTokenizer
 
 OPENAI_KEY = os.getenv("OPENAI_API_KEY")
 HF_TOKEN = os.getenv("HF_TOKEN")
-TOKENIZER = AutoTokenizer.from_pretrained(os.getenv("HF_MODEL"))
+#TOKENIZER = AutoTokenizer.from_pretrained(os.getenv("HF_MODEL"))
 
-HF_CLIENT = InferenceClient(
-    os.getenv("HF_MODEL"),
-    token=HF_TOKEN
-)
+#HF_CLIENT = InferenceClient(
+#    os.getenv("HF_MODEL"),
+#    token=HF_TOKEN
+#)
 OAI_CLIENT = openai.Client(api_key=OPENAI_KEY)
 
 HF_GENERATE_KWARGS = {
@@ -115,7 +115,6 @@ def generate_openai(prompt: str, history: str, model_name: str) -> Generator[str
 
     try:
         stream = OAI_CLIENT.chat.completions.create(
-            #model=os.getenv("OPENAI_MODEL"),
             model = model_name,
             messages=formatted_prompt,
             **OAI_GENERATE_KWARGS,