Update app.py

app.py

@@ -13,8 +13,11 @@ from threading import Thread
 from FlagEmbedding import BGEM3FlagModel
 from sklearn.metrics.pairwise import cosine_similarity
 
+from transformers import AutoModelForSequenceClassification
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+#Importing the embedding model
 embedding_model = BGEM3FlagModel('BAAI/bge-m3',  
                        use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
 
@@ -22,20 +25,44 @@ embeddings = np.load("embeddings_albert_tchap.npy")
 embeddings_data = pd.read_json("embeddings_albert_tchap.json")
 embeddings_text = embeddings_data["text_with_context"].tolist()
 
-# Define the device
-temperature=0.2
-max_new_tokens=1000
-top_p=0.92
-repetition_penalty=1.7
+#Importing the classifier/router (deberta)
+classifier_model = AutoModelForSequenceClassification.from_pretrained("AgentPublic/chatrag-deberta")
+tokenizer = AutoTokenizer.from_pretrained("AgentPublic/chatrag-deberta")
 
+#Importing the actual generative LLM (llama-based)
 model_name = "Pclanglais/Tchap"
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
 model = model.to('cuda:0')
 
 system_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nTu es Albert, l'agent conversationnel des services publics qui peut décrire des documents de référence ou aider à des tâches de rédaction<|eot_id|>"
 
+#Function to guess whether we use the RAG or not.
+def classification_chatrag(query):
+  encoding = tokenizer(query, return_tensors="pt")
+  encoding = {k: v.to(model_classifier.device) for k,v in encoding.items()}
+
+  outputs = model_classifier(**encoding)
+
+  logits = outputs.logits
+  logits.shape
+
+  # apply sigmoid + threshold
+  sigmoid = torch.nn.Sigmoid()
+  probs = sigmoid(logits.squeeze().cpu())
+  predictions = np.zeros(probs.shape)
+
+  # Extract the float value from the tensor
+  float_value = round(probs.item()*100)
+
+  if float_value > 50:
+    status = True
+    print("We activate RAG")
+  else:
+    status = False
+    print("We remove RAG")
+  return status
+
 #Vector search over the database
 def vector_search(sentence_query):
 
@@ -71,9 +98,12 @@ class StopOnTokens(StoppingCriteria):
 def predict(message, history):
 
     global source_text
+    global assess_rag
     #For now, we only query the vector database once, at the start.
     if len(history) == 0:
-        source_text = vector_search(message)
+        assess_rag = classification_chatrag(message)
+        if assess_rag:
+            source_text = vector_search(message)
     
     history_transformer_format = history + [[message, ""]]
 
@@ -87,7 +117,10 @@ def predict(message, history):
 
         #Once we target the ongoing post we add the source.
         if id_message == total_message:
-            question = "<|start_header_id|>user<|end_header_id|>\n\n"+ item[0] + "\n\n### Source ###\n" + source_text
+            if assess_rag:
+                question = "<|start_header_id|>user<|end_header_id|>\n\n"+ item[0] + "\n\n### Source ###\n" + source_text
+            else:
+                question = "<|start_header_id|>user<|end_header_id|>\n\n"+ item[0]
         else:
             question = "<|start_header_id|>user<|end_header_id|>\n\n"+ item[0]
         answer = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"+item[1]