app.py

import gradio as gr
from transformers import pipeline
from huggingface_hub import InferenceClient, login, snapshot_download
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
import os
import pandas as pd


"""
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)

#model = "meta-llama/Llama-3.2-1B-Instruct"
#model = "google/mt5-small"
model = "mistralai/Mistral-7B-Instruct-v0.3"

client = InferenceClient(model)

folder = snapshot_download(repo_id="umaiku/faiss_index", repo_type="dataset", local_dir=os.getcwd())

embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-small")

vector_db = FAISS.load_local("faiss_index_8k", embeddings, allow_dangerous_deserialization=True)

df = pd.read_csv("faiss_index/bger_cedh_db 1954-2024.csv")

"""
You are an assistant in Swiss Jurisprudence law.
Please answer the user in the same language that he used in his question using the following given context, not prior or other knowledge.
If no relevant cases were retrieved 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 and do not give any links that are not provided in the context.
Otherwise, if relevant cases were found, start by summarizing these cases in the user's question's language and reference the sources, including the source, urls and dates.
"""

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    score,
):
    messages = [{"role": "system", "content": system_message}]

    print(system_message)
    
    retriever = vector_db.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": score})
    documents = retriever.invoke(message)

    spacer = " \n"
    context = ""
    print(len(documents))

    for doc in documents:
        case_text = df[df["case_url"] == doc.metadata["case_url"]].case_text.values[0]

        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
    
    
    message = f"""
The user is asking you the following question: {message}
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: Always answer the user using the language used in his question: {message}
    """
    
    print(message)

#    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})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


"""
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="", 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.7, 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, maximum=1, value=0.7, step=0.1, label="Score Threshold"),
    ],
    description="# 📜 ALexI: Artificial Legal Intelligence for Swiss Jurisprudence",
)


if __name__ == "__main__":
    demo.launch(debug=True)