Update app.py

Back to version before ChatGPT

app.py

@@ -1,125 +1,19 @@
-import os
 import gradio as gr
-import pandas as pd
-from datetime import datetime
-
 from transformers import pipeline
 from huggingface_hub import InferenceClient, login, snapshot_download
-
 from langchain_community.vectorstores import FAISS, DistanceStrategy
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_core.vectorstores import VectorStore
+import os
+import pandas as pd
+from datetime import datetime
 
 from smolagents import Tool, HfApiModel, ToolCallingAgent
+from langchain_core.vectorstores import VectorStore
 
 
-# -------- Helpers & Compatibility --------
-
-def _warn_token():
-    hf_token = os.getenv("TOKEN") or os.getenv("HF_TOKEN")
-    if not hf_token:
-        print("[WARN] No HF token found in env (TOKEN or HF_TOKEN). Private models/endpoints may fail.")
-        return None
-    return hf_token
-
-def _login_hf():
-    token = _warn_token()
-    if token:
-        try:
-            login(token=token)
-        except TypeError:
-            # older huggingface_hub accepted positional
-            login(token)
-
-def _stream_chat(client: InferenceClient, messages, max_tokens: int, temperature: float, top_p: float):
-    """
-    Try new OpenAI-style streaming first, then older `chat_completion`, then fall back to text_generation.
-    Yields string chunks.
-    """
-    # 1) New: client.chat.completions.create(..., stream=True)
-    try:
-        chat = client.chat.completions.create(
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True,
-        )
-        for chunk in chat:
-            # choices[0].delta.content may be None in some chunks
-            if chunk and getattr(chunk, "choices", None):
-                delta = chunk.choices[0].delta
-                if delta and getattr(delta, "content", None):
-                    yield delta.content
-        return
-    except Exception as e_new:
-        # print for debug, but continue to fallback
-        print("[INFO] OpenAI-style chat.completions streaming not available:", repr(e_new))
-
-    # 2) Old: client.chat_completion(..., stream=True)
-    try:
-        old_stream = client.chat_completion(
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True,
-        )
-        # Old stream objects sometimes have .choices[0].delta.content, sometimes just .token
-        for chunk in old_stream:
-            text = None
-            try:
-                text = chunk.choices[0].delta.content  # may exist
-            except Exception:
-                pass
-            if not text:
-                # try common fallbacks
-                text = getattr(chunk, "token", None) or getattr(chunk, "text", None)
-            if text:
-                yield text
-        return
-    except Exception as e_old:
-        print("[INFO] Legacy chat_completion streaming not available:", repr(e_old))
-
-    # 3) Fallback: plain text_generation with a single concatenated prompt (no messages)
-    # The last user message should be the final prompt.
-    try:
-        final_prompt = ""
-        for m in messages:
-            role = m.get("role", "user")
-            content = m.get("content", "")
-            # simple role-tagged concat
-            final_prompt += f"{role.upper()}: {content}\n"
-        gen_stream = client.text_generation(
-            final_prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True,
-            return_full_text=False,
-        )
-        for piece in gen_stream:
-            # piece may be string or an object with .token/.generated_text
-            if isinstance(piece, str):
-                yield piece
-            else:
-                text = getattr(piece, "token", None) or getattr(piece, "generated_text", None)
-                if text:
-                    yield text
-        return
-    except Exception as e_gen:
-        print("[ERROR] All HF streaming methods failed:", repr(e_gen))
-        yield "\n[Error] Unable to stream from the inference endpoint. Check model name, token, and HF API version.\n"
-
-
-# -------- Data / Vector Store --------
-
 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."
-    )
+    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",
@@ -128,169 +22,141 @@ class RetrieverTool(Tool):
     }
     output_type = "string"
 
-    def __init__(self, vectordb: VectorStore, df: pd.DataFrame, **kwargs):
+    def __init__(self, vectordb: VectorStore, **kwargs):
         super().__init__(**kwargs)
         self.vectordb = vectordb
-        self.df = df
 
     def forward(self, query: str) -> str:
         assert isinstance(query, str), "Your search query must be a string"
-        try:
-            docs = self.vectordb.similarity_search(query, k=7)
-        except Exception as e:
-            return f"[Retriever error] {e}"
+
+        docs = self.vectordb.similarity_search(
+            query,
+            k=7,
+        )
 
         spacer = " \n"
         context = ""
         nb_char = 100
-
+        
         for doc in docs:
-            # Safe metadata access
-            meta = getattr(doc, "metadata", {}) or {}
-            case_ref = str(meta.get("case_ref", "") or "")
-            case_nb = str(meta.get("case_nb", "") or "")
-            case_date = str(meta.get("case_date", "") or "")
-            case_url = str(meta.get("case_url", "") or "")
-
-            # Try to find a surrounding extract from the master text
-            case_text_summary = ""
-            if case_url:
-                try:
-                    rows = self.df[self.df["case_url"] == case_url]
-                    if not rows.empty:
-                        case_text = str(rows.iloc[0]["case_text"])
-                        idx = case_text.find(doc.page_content)
-                        if idx >= 0:
-                            start = max(0, idx - nb_char)
-                            end = min(len(case_text), idx + len(doc.page_content) + nb_char)
-                            case_text_summary = case_text[start:end]
-                except Exception as e:
-                    # If anything goes wrong, fall back to page_content
-                    case_text_summary = doc.page_content
-
-            if not case_text_summary:
-                case_text_summary = doc.page_content
-
+            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: " + (case_ref + " " + case_nb).strip() + spacer
-            source_name = "Swiss Federal Court" if case_ref == "ATF" else "European Court of Human Rights"
-            context += "# Case source: " + source_name + spacer
-            context += "# Case date: " + case_date + spacer
-            context += "# Case url: " + case_url + 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
 
+        return "\nRetrieved documents:\n" + context
 
-# -------- Init HF / Model / Index --------
 
-_login_hf()
+"""
+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
+"""
+HF_TOKEN=os.getenv('TOKEN')
+login(HF_TOKEN)
 
-# Choose your model
-MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
-# MODEL_ID = "swiss-ai/Apertus-8B-Instruct-2509"
+model = "meta-llama/Meta-Llama-3-8B-Instruct"
+#model = "swiss-ai/Apertus-8B-Instruct-2509"
 
-client = InferenceClient(MODEL_ID)
+client = InferenceClient(model)
 
-# Pull the FAISS dataset snapshot and derive the index path
 folder = snapshot_download(repo_id="umaiku/faiss_index", repo_type="dataset", local_dir=os.getcwd())
-index_dir = os.path.join(folder, "faiss_index_mpnet_cos")
-if not os.path.isdir(index_dir):
-    # Fallback: try current working directory if you’ve manually placed the index there
-    alt = os.path.join(os.getcwd(), "faiss_index_mpnet_cos")
-    if os.path.isdir(alt):
-        index_dir = alt
-    else:
-        print(f"[WARN] Could not find FAISS index directory at {index_dir} or {alt}. Check your dataset contents.")
 
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2")
 
-# Load FAISS (COSINE distance)
-vector_db = FAISS.load_local(
-    index_dir,
-    embeddings,
-    allow_dangerous_deserialization=True,
-    distance_strategy=DistanceStrategy.COSINE,
-)
-
-# Load your case dataframe
-CSV_PATH = os.path.join(folder, "bger_cedh_db 1954-2024.csv")
-if not os.path.isfile(CSV_PATH):
-    # also try local if you keep it next to the script
-    CSV_PATH = "bger_cedh_db 1954-2024.csv"
-df = pd.read_csv(CSV_PATH)
+vector_db = FAISS.load_local("faiss_index_mpnet_cos", embeddings, allow_dangerous_deserialization=True, distance_strategy=DistanceStrategy.COSINE)
 
-retriever_tool = RetrieverTool(vector_db, df)
-agent = ToolCallingAgent(tools=[retriever_tool], model=HfApiModel(MODEL_ID))  # Not used directly, but kept if you expand.
+df = pd.read_csv("bger_cedh_db 1954-2024.csv")
 
+retriever_tool = RetrieverTool(vector_db)
+agent = ToolCallingAgent(tools=[retriever_tool], model=HfApiModel(model))
 
-# -------- Chat callback --------
+def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p, score,):
 
-def respond(
-    user_message: str,
-    history: list[tuple[str, str]],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-    score_threshold: float,
-):
     print(datetime.now())
-    print("[User]", user_message)
+    context = retriever_tool(message)
+    
+    print(message)
 
-    context = retriever_tool(user_message)
+#    is_law = client.text_generation(f"""Given the user question below, classify it as either being about "Law" or "Other".
+#Do NOT respond with more than one word.
+#Question:
+#{message}""")
 
-    # Build the RAG prompt
-    prompt = f"""Given the question and supporting documents below, give a comprehensive answer to the question.
+#    print(is_law)
+
+    if True: #is_law.lower() != "other":    
+        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 relevant to the question and 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!
 Always give url of the sources at the end and only answer in the language the question is asked.
-
+    
 Question:
-{user_message}
-
+{message}
+    
 {context}
 """
-
+    else:
+        prompt = f"""A user wrote the following message, please answer him to best of your knowledge in the language of his message:
+{message}"""
+    
     messages = [{"role": "system", "content": system_message}]
 
-    # Rehydrate prior turns (user, assistant)
-    for u, a in history:
-        if u:
-            messages.append({"role": "user", "content": u})
-        if a:
-            messages.append({"role": "assistant", "content": a})
+    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": prompt})
 
-    response_accum = ""
-    for chunk_text in _stream_chat(
-        client,
-        messages=messages,
+    response = ""
+
+
+    for message in client.chat_completion(
+        messages,
         max_tokens=max_tokens,
+        stream=True,
         temperature=temperature,
         top_p=top_p,
     ):
-        if chunk_text:
-            response_accum += chunk_text
-            yield response_accum
+        token = message.choices[0].delta.content
+        
+        response += token
+        yield response
 
 
-# -------- Gradio UI --------
-
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are assisting a jurist or a lawyer in finding relevant Swiss Jurisprudence cases to their question.", label="System message"),
+        gr.Textbox(value="You are assisting a jurist or a layer in finding relevant Swiss Jurisprudence cases to their question.", 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.1, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
         gr.Slider(minimum=0, maximum=1, value=0.75, step=0.05, label="Score Threshold"),
     ],
     description="# 📜 ALexI: Artificial Legal Intelligence for Swiss Jurisprudence",
 )
 
+
 if __name__ == "__main__":
     print("Ready!")
-    demo.launch(debug=True)
+    demo.launch(debug=True)