Update app.py

Translate prompt and extract longer context chunks

app.py

@@ -41,28 +41,53 @@ def respond(
 
     print(datetime.now())
     print(system_message)
+
+    prompt_en = client.text_generation("Improve or translate the following user's prompt to English giving only the new prompt without explanations or additional text: " + message)
+    prompt_de = client.text_generation("Improve or translate the following user's prompt to German giving only the new prompt without explanations or additional text: " + message)
+    prompt_fr = client.text_generation("Improve or translate the following user's prompt to French giving only the new prompt without explanations or additional text: " + message)
+    prompt_it = client.text_generation("Improve or translate the following user's prompt to Italian giving only the new prompt without explanations or additional text: " + message)
+
+    print(prompt_en)
+    print(prompt_de)
+    print(prompt_fr)
+    print(prompt_it)
+    
     
 #    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="similarity", search_kwargs={"k": 10})
 #    retriever = vector_db.as_retriever(search_type="mmr")
-    documents = retriever.invoke(message)
+#    documents = retriever.invoke(message)
+
+    documents_en = vector_db.similarity_search_with_score(prompt_en, k=4)
+    documents_de = vector_db.similarity_search_with_score(prompt_de, k=4)
+    documents_fr = vector_db.similarity_search_with_score(prompt_fr, k=4)
+    documents_it = vector_db.similarity_search_with_score(prompt_it, k=4)
 
+    documents = documents_en + documents_de + documents_fr + documents_it
+
+    documents = sorted(documents, key=lambda x: x[1])[:4]
+    
     spacer = " \n"
     context = ""
+    nb_char = 2000
 
     #print(message)
-    print(len(documents))
+    print(f"* Documents found: {len(documents)}")
 
     for doc in documents:
-        #case_text = df[df["case_url"] == doc.metadata["case_url"]].case_text.values[0]
-
+        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.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 text: " + case_text[:8000] + 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)
@@ -114,7 +139,7 @@ demo = gr.ChatInterface(
     additional_inputs=[
         gr.Textbox(value="You are an assistant in Swiss Jurisprudence cases.", label="System message"),
         gr.Slider(minimum=1, maximum=24000, value=5000, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.1, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,