Update app.py

app.py

@@ -1,12 +1,63 @@
 import gradio as gr
 from transformers import pipeline
 from huggingface_hub import InferenceClient, login, snapshot_download
-from langchain_community.vectorstores import FAISS
+from langchain_community.vectorstores import FAISS, DistanceStrategy
 from langchain_huggingface import HuggingFaceEmbeddings
 import os
 import pandas as pd
 from datetime import datetime
 
+from smolagents import Tool, HfApiModel, ToolCallingAgent
+from langchain_core.vectorstores import VectorStore
+
+
+class RetrieverTool(Tool):
+    name = "retriever"
+    description = "Using semantic similarity in German, French, English and Italian, retrieves some documents from the knowledge base that have the closest embeddings to the input query."
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The query to perform. This should be semantically close to your target documents. Use the affirmative form rather than a question.",
+        }
+    }
+    output_type = "string"
+
+    def __init__(self, vectordb: VectorStore, **kwargs):
+        super().__init__(**kwargs)
+        self.vectordb = vectordb
+
+    def forward(self, query: str) -> str:
+        assert isinstance(query, str), "Your search query must be a string"
+
+        docs = self.vectordb.similarity_search(
+            query,
+            k=7,
+        )
+
+        df = pd.read_csv("bger_cedh_db 1954-2024.csv")
+        
+        spacer = " \n"
+        context = ""
+        nb_char = 100
+        
+        for doc in docs:
+            case_text = df[df["case_url"] == doc.metadata["case_url"]].case_text.values[0]
+            index = case_text.find(doc.page_content)
+            start = max(0, index - nb_char)
+            end = min(len(case_text), index + len(doc.page_content) + nb_char)
+            case_text_summary = case_text[start:end]
+            
+            context += "#######" + spacer
+            context += "# Case number: " + doc.metadata["case_ref"] + " " + doc.metadata["case_nb"] + spacer
+            context += "# Case source: " + ("Swiss Federal Court" if doc.metadata["case_ref"] == "ATF" else "European Court of Human Rights") + spacer
+            context += "# Case date: " + doc.metadata["case_date"] + spacer
+            context += "# Case url: " + doc.metadata["case_url"] + spacer
+            #context += "# Case text: " + doc.page_content + spacer
+            context += "# Case extract: " + case_text_summary + spacer
+
+
+        return "\nRetrieved documents:\n" + context
+
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
@@ -14,9 +65,7 @@ For more information on `huggingface_hub` Inference API support, please check th
 HF_TOKEN=os.getenv('TOKEN')
 login(HF_TOKEN)
 
-#model = "meta-llama/Llama-3.2-1B-Instruct"
-#model = "google/mt5-small"
-model = "mistralai/Mistral-7B-Instruct-v0.3"
+model = "meta-llama/Meta-Llama-3-8B-Instruct"
 
 client = InferenceClient(model)
 
@@ -24,98 +73,36 @@ folder = snapshot_download(repo_id="umaiku/faiss_index", repo_type="dataset", lo
 
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2")
 
-vector_db = FAISS.load_local("faiss_index_mpnet_cos", embeddings, allow_dangerous_deserialization=True)
+vector_db = FAISS.load_local("faiss_index_mpnet_cos", embeddings, allow_dangerous_deserialization=True, distance_strategy=DistanceStrategy.COSINE)
 
 df = pd.read_csv("faiss_index/bger_cedh_db 1954-2024.csv")
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    score,
-):
-    #messages = [{"role": "system", "content": system_message}]
+retriever_tool = RetrieverTool(vector_db)
+agent = ToolCallingAgent(tools=[retriever_tool], model=HfApiModel(model))
 
-    print(datetime.now())
-    print(system_message)
+def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p, score,):
 
-    prompt_template = "Improve or translate the following user's prompt to {language} giving only the new prompt\
-    without explanations or additional text and if you can't improve it, just return the same prompt, do not extrapolate: "
+    print(datetime.now())
+    context = retriever_tool(question)
     
-    prompt_en = client.text_generation(prompt_template.format(language="English") + message)
-    prompt_de = client.text_generation(prompt_template.format(language="German") + message)
-    prompt_fr = client.text_generation(prompt_template.format(language="French") + message)
-    prompt_it = client.text_generation(prompt_template.format(language="Italian") + message)
-
-#    retriever = vector_db.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": score, "k": 10})
-#    retriever = vector_db.as_retriever(search_type="similarity", search_kwargs={"k": 10})
-#    retriever = vector_db.as_retriever(search_type="mmr")
-#    documents = retriever.invoke(message)
-
-    documents_en = vector_db.similarity_search_with_score(prompt_en, k=4)
-    print(prompt_en)
+    prompt = f"""Given the question and supporting documents below, give a comprehensive answer to the question.
+    Respond only to the question asked, response should be concise and relevant to the question and answer in the same language as the question.
+    Provide the number of the source document when relevant, as well as the link to the document.
+    If you cannot find information, do not give up and try calling your retriever again with different arguments!
     
-    documents_de = vector_db.similarity_search_with_score(prompt_de, k=4)
-    print(prompt_de)
-    documents_fr = vector_db.similarity_search_with_score(prompt_fr, k=4)
-    print(prompt_fr)
-
-    documents_it = vector_db.similarity_search_with_score(prompt_it, k=4)
-    print(prompt_it)
-
-    documents = documents_en + documents_de + documents_fr + documents_it
-
-    documents = sorted(documents, key=lambda x: x[1])[:4]
+    Question:
+    {question}
     
-    spacer = " \n"
-    context = ""
-    nb_char = 2000
-
-    #print(message)
-    print(f"* Documents found: {len(documents)}")
-
-    for doc in documents:
-        case_text = df[df["case_url"] == doc[0].metadata["case_url"]].case_text.values[0]
-        index = case_text.find(doc[0].page_content)
-        start = max(0, index - nb_char)
-        end = min(len(case_text), index + len(doc[0].page_content) + nb_char)
-        case_text_summary = case_text[start:end]
-        
-        context += "#######" + spacer
-        context += "# Case number: " + doc[0].metadata["case_nb"] + spacer
-        context += "# Case source: " + ("Swiss Federal Court" if doc[0].metadata["case_ref"] == "ATF" else "European Court of Human Rights") + spacer
-        context += "# Case date: " + doc[0].metadata["case_date"] + spacer
-        context += "# Case url: " + doc[0].metadata["case_url"] + spacer
-        #context += "# Case text: " + doc[0].page_content + spacer
-        context += "Case extract: " + case_text_summary + spacer
-
-        #print("# Case number: " + doc.metadata["case_nb"] + spacer)
-        #print("# Case url: " + doc.metadata["case_url"] + spacer)
-   
-    system_message += f"""A user is asking you the following question: {message}
-Please answer the user in the same language that he used in his question using ONLY the following given context not any prior knowledge or information found on the internet.
-# Context:
-The following case extracts have been found in either Swiss Federal Court or European Court of Human Rights cases and could fit the question:
-{context}
-# Task:
-If the retrieved context is not relevant cases or the issue has not been addressed within the context, just say "I can't find enough relevant information".
-Don't make up an answer or give irrelevant information not requested by the user.
-Otherwise, if relevant cases were found, answer in the user's question's language using the context that you found relevant and reference the sources, including the urls and dates.
-# Instructions:
-Always answer the user using the language used in his question: {message}
-"""
-
-    print(system_message)
-    messages = [{"role": "system", "content": system_message}]
+    {context}
+    """
     
-#    for val in history:
-#        if val[0]:
-#            messages.append({"role": "user", "content": val[0]})
-#        if val[1]:
-#            messages.append({"role": "assistant", "content": val[1]})
+    messages = [{"role": "user", "content": prompt}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
 
     messages.append({"role": "user", "content": message})
 
@@ -129,6 +116,8 @@ Always answer the user using the language used in his question: {message}
         top_p=top_p,
     ):
         token = message.choices[0].delta.content
+        
+#        answer = client.chat_completion(messages, temperature=0.1).choices[0].message.content
 
         response += token
         yield response