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

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"""
A user is asking you the following question: {message}
Please answer the user in the same language that he used in his question using the following given context not prior knowledge.
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, start by summarizing these cases in the user's question's language and reference the sources, including the urls and dates.
Instructions:
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="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.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)