Update app.py

app.py

@@ -1,16 +1,25 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-import torch
+from datasets import load_dataset
+
 import os
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+import torch
+from threading import Thread
+from sentence_transformers import SentenceTransformer
+from datasets import load_dataset
+import time
 
-# Check CUDA availability
-print("CUDA available:", torch.cuda.is_available())
-print("CUDA version:", torch.version.cuda)
+token = os.environ["HF_TOKEN"]
+ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
 
-# Load the Hugging Face token from secrets
-token = os.environ.get("HF_TOKEN")
+dataset = load_dataset("not-lain/wikipedia",revision = "embedded")
 
-model_id = "google/shieldgemma-2b"
+data = dataset["train"]
+data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset
+
+
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 
 # use quantization to lower GPU usage
 bnb_config = BitsAndBytesConfig(
@@ -25,109 +34,104 @@ model = AutoModelForCausalLM.from_pretrained(
     quantization_config=bnb_config,
     token=token
 )
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
 
-# Función para generar el prompt dependiendo del idioma seleccionado
-def generar_prompt(message, tipo_clasificacion, idioma):
-    if idioma == "Español":
-        politicas = {
-            "Contenido Peligroso": """
-            * "No al contenido peligroso": El mensaje no debe contener ni buscar la generación de contenido que cause daño a uno mismo o a otros (por ejemplo, acceder o fabricar armas de fuego o explosivos, promoción del terrorismo, instrucciones para el suicidio).
-            """,
-            "Acoso": """
-            * "No al acoso": El mensaje no debe contener ni buscar la generación de contenido que sea malicioso, intimidante, que fomente el bullying o abusivo dirigido a otra persona (por ejemplo, amenazas físicas, negación de eventos trágicos, menosprecio a víctimas de violencia).
-            """,
-            "Discurso de Odio": """
-            * "No al discurso de odio": El mensaje no debe contener ni buscar la generación de contenido que ataque la identidad o atributos protegidos (por ejemplo, insultos raciales, promoción de la discriminación, llamados a la violencia contra grupos protegidos) o que deshumanice, menosprecie o vilipendie por motivos de raza, etnia, religión, discapacidad, edad, nacionalidad, estatus de veterano, orientación sexual, género, identidad de género, casta, u otro estatus protegido.
-            """,
-            "Contenido Sexual": """
-            * "No al contenido sexual explícito": El mensaje no debe contener ni buscar la generación de contenido que haga referencia a actos sexuales u otro contenido lascivo (por ejemplo, descripciones gráficas sexuales, contenido destinado a causar excitación). Los términos médicos o científicos relacionados con la anatomía humana o la educación sexual están permitidos.
-            """
-        }
-    else:  # English
-        politicas = {
-            "Dangerous Content": """
-            * "No Dangerous Content": The prompt shall not contain or seek generation of content that harms oneself and/or others (e.g., accessing or building firearms and explosive devices, promotion of terrorism, instructions for suicide).
-            """,
-            "Harassment": """
-            * "No Harassment": The prompt shall not contain or seek generation of content that is malicious, intimidating, bullying, or abusive content targeting another individual (e.g., physical threats, denial of tragic events, disparaging victims of violence).
-            """,
-            "Hate Speech": """
-            * "No Hate Speech": The prompt shall not contain or seek generation of content that targets identity and/or protected attributes (e.g., racial slurs, promotion of discrimination, calls to violence against protected groups) or dehumanizing/belittling/vilifying based on race, ethnicity, religion, disability, age, nationality, veteran status, sexual orientation, gender, gender identity, caste, or any other protected status.
-            """,
-            "Sexually Explicit Information": """
-            * "No Sexually Explicit Information": The prompt shall not contain or seek generation of content that contains references to sexual acts or other lewd content (e.g., sexually graphic descriptions, content aimed at causing arousal). Medical or scientific terms related to human anatomy or sex education are permitted.
-            """
-        }
-
-    # Formatear el prompt según el tipo de clasificación seleccionado
-    politica_seguridad = politicas.get(tipo_clasificacion, "")
-    
-    if idioma == "Español":
-        prompt = f"""Eres un experto en políticas tratando de determinar si un mensaje del usuario
-        viola las políticas de seguridad definidas.
-        <inicio_turno>
-        Pregunta del usuario: {message.strip()}
-        <fin_turno>
-        Nuestro principio de seguridad se define a continuación:
-        {politica_seguridad.strip()}
-        ¿La pregunta del usuario viola el principio anterior? Tu respuesta debe comenzar con
-        'Sí' o 'No'. Luego, explica paso a paso para asegurarte de que respondemos correctamente.
-        """
-    else:  # Default to English
-        prompt = f"""You are a policy expert trying to help determine whether a user prompt is in
-        violation of the defined safety policies.
-        <start_of_turn>
-        Human Question: {message.strip()}
-        <end_of_turn>
-        Our safety principle is defined below:
-        {politica_seguridad.strip()}
-        Does the human question violate the above principle? Your answer must start with
-        'Yes' or 'No'. Then walk through step by step to be sure we answer correctly.
-        """
-    
-    return prompt
-
-@spaces.GPU(duration=150)
-# Función para procesar la respuesta y clasificar según las políticas de seguridad
-def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p, language, harm_type):
-    prompt = generar_prompt(message, harm_type, language)
+SYS_PROMPT = """You are an assistant for answering questions.
+You are given the extracted parts of a long document and a question. Provide a conversational answer.
+If you don't know the answer, just say "I do not know." Don't make up an answer."""
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)["input_ids"]
 
-    with torch.no_grad():
-        logits = model(**inputs).logits
 
-    # Extraer los logits para los tokens 'Yes'/'No' en inglés o 'Sí'/'No' en español
-    vocab = tokenizer.get_vocab()
-    if language == "Español":
-        selected_logits = logits[0, -1, [vocab['Sí'], vocab['No']]]
-    else:
-        selected_logits = logits[0, -1, [vocab['Yes'], vocab['No']]]
+def search(query: str, k: int = 3 ):
+    """a function that embeds a new query and returns the most probable results"""
+    embedded_query = ST.encode(query) # embed new query
+    scores, retrieved_examples = data.get_nearest_examples( # retrieve results
+        "embeddings", embedded_query, # compare our new embedded query with the dataset embeddings
+        k=k # get only top k results
+    )
+    return scores, retrieved_examples
 
-    # Convertir los logits en una probabilidad con softmax
-    probabilities = softmax(selected_logits, dim=0)
+def format_prompt(prompt,retrieved_documents,k):
+    """using the retrieved documents we will prompt the model to generate our responses"""
+    PROMPT = f"Question:{prompt}\nContext:"
+    for idx in range(k) :
+        PROMPT+= f"{retrieved_documents['text'][idx]}\n"
+    return PROMPT
 
-    # Devolver la probabilidad de 'Sí'/'Yes' y la respuesta generada
-    score_yes = probabilities[0].item()
 
-    if language == "Español":
-        response = f"Puntuación para 'Sí' (violación): {score_yes:.4f}"
-    else:
-        response = f"Score for 'Yes' (violation): {score_yes:.4f}"
+@spaces.GPU(duration=150)
+def talk(prompt,history):
+    k = 1 # number of retrieved documents
+    scores , retrieved_documents = search(prompt, k)
+    formatted_prompt = format_prompt(prompt,retrieved_documents,k)
+    formatted_prompt = formatted_prompt[:2000] # to avoid GPU OOM
+    messages = [{"role":"system","content":SYS_PROMPT},{"role":"user","content":formatted_prompt}]
+    # tell the model to generate
+    input_ids = tokenizer.apply_chat_template(
+      messages,
+      add_generation_prompt=True,
+      return_tensors="pt"
+    ).to(model.device)
+    outputs = model.generate(
+      input_ids,
+      max_new_tokens=1024,
+      eos_token_id=terminators,
+      do_sample=True,
+      temperature=0.6,
+      top_p=0.9,
+    )
+    streamer = TextIteratorStreamer(
+            tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+        )
+    generate_kwargs = dict(
+        input_ids= input_ids,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.95,
+        temperature=0.75,
+        eos_token_id=terminators,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        print(outputs)
+        yield "".join(outputs)
+
+
+TITLE = "# RAG"
+
+DESCRIPTION = """
+A rag pipeline with a chatbot feature
+Resources used to build this project :
+* embedding model : https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1
+* dataset : https://huggingface.co/datasets/not-lain/wikipedia
+* faiss docs : https://huggingface.co/docs/datasets/v2.18.0/en/package_reference/main_classes#datasets.Dataset.add_faiss_index 
+* chatbot : https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct
+* Full documentation : https://huggingface.co/blog/not-lain/rag-chatbot-using-llama3 
+"""
 
-    return response
 
-# Crear la interfaz de Gradio con selección de idioma y tipo de contenido
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, 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.Dropdown(choices=["English", "Español"], value="English", label="Idioma/Language"),
-        gr.Dropdown(choices=["Dangerous Content", "Harassment", "Hate Speech", "Sexually Explicit Information"], value="Harassment", label="Harm Type")
-    ],
+    fn=talk,
+    chatbot=gr.Chatbot(
+        show_label=True,
+        show_share_button=True,
+        show_copy_button=True,
+        likeable=True,
+        layout="bubble",
+        bubble_full_width=False,
+    ),
+    theme="Soft",
+    examples=[["what's anarchy ? "]],
+    title=TITLE,
+    description=DESCRIPTION,
+    
 )
-
 demo.launch(debug=True)