Update app.py

app.py

@@ -1,93 +1,113 @@
 import gradio as gr
 from datasets import load_dataset
-
 import os
-import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
-import torch
 from threading import Thread
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from sentence_transformers import SentenceTransformer
-from datasets import load_dataset
-import time
 
+# Configuración y token
 token = os.environ["HF_TOKEN"]
-ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
-
-dataset = load_dataset("not-lain/wikipedia",revision = "embedded")
 
+# Modelo de embeddings y dataset para la fase de RAG
+ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
+dataset = load_dataset("not-lain/wikipedia", revision="embedded")
 data = dataset["train"]
-data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset
+data = data.add_faiss_index("embeddings")  # columna con los embeddings
 
-
-model_id = "meta-llama/Llama-3.2-3B-Instruct"
-
-# use quantization to lower GPU usage
+# Configuración de quantization
 bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    load_in_4bit=True, 
+    bnb_4bit_use_double_quant=True, 
+    bnb_4bit_quant_type="nf4", 
+    bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_id,token=token)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    quantization_config=bnb_config,
-    token=token
+# Diccionario con modelos disponibles
+# Puedes agregar tantos modelos como desees, con su respectiva configuración
+MODELOS = {
+    "Llama-3.2-3B-Instruct": "meta-llama/Llama-3.2-3B-Instruct",
+    "Otro Modelo": "ruta/o/identificador-del-otro-modelo"
+}
+
+# Opcional: Diccionario para almacenar modelos cargados y evitar recargas
+MODELOS_CARGADOS = {}
+
+# Función para obtener (o cargar) el modelo y tokenizer según la selección
+def get_model(selected_model):
+    model_id = MODELOS[selected_model]
+    if selected_model not in MODELOS_CARGADOS:
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            quantization_config=bnb_config,
+            token=token
+        )
+        MODELOS_CARGADOS[selected_model] = (model, tokenizer)
+    return MODELOS_CARGADOS[selected_model]
+
+# Tokens de finalización basados en el tokenizer del modelo, se actualizarán luego
+def get_terminators(tokenizer):
+    return [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
+
+# Prompt del sistema (podrías parametrizarlo si es necesario)
+SYS_PROMPT = (
+    "Tu tarea es analizar un listado de unidades de competencia y devolver únicamente aquellas relacionadas "
+    "con la profesión de {texto_usuario}. Debes buscar palabras clave o términos relacionados con la profesión "
+    "en los nombres de las unidades de competencia. If you don't know the answer, just say 'I do not know.' "
+    "Don't make up an answer."
 )
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
-
-SYS_PROMPT = """Tu tarea es analizar un listado de unidades de competencia y devolver únicamente aquellas relacionadas con la profesión de {texto_usuario}.
-Debes buscar palabras clave o términos relacionados con la profesión en los nombres de las unidades de competencia.
-If you don't know the answer, just say "I do not know." Don't make up an answer."""
-
 
-
-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
+def search(query: str, k: int = 3):
+    """Embebe la consulta y retorna los resultados más probables."""
+    embedded_query = ST.encode(query)  # embed de la nueva consulta
+    scores, retrieved_examples = data.get_nearest_examples(
+        "embeddings",
+        embedded_query,
+        k=k
     )
     return scores, retrieved_examples
 
-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"
+def format_prompt(prompt, retrieved_documents, k):
+    """Construye el prompt a partir de los documentos recuperados."""
+    PROMPT = f"Question: {prompt}\nContext:\n"
+    for idx in range(k):
+        PROMPT += f"{retrieved_documents['text'][idx]}\n"
     return PROMPT
 
-
-@spaces.GPU
-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
+def talk(prompt, selected_model, history=[]):
+    # Obtiene (o carga) el modelo y tokenizer basado en la selección
+    model, tokenizer = get_model(selected_model)
+    terminators = get_terminators(tokenizer)
+    
+    # Número de documentos a recuperar
+    k = 1
+    scores, retrieved_documents = search(prompt, k)
+    formatted_prompt = format_prompt(prompt, retrieved_documents, k)
+    formatted_prompt = formatted_prompt[:2000]  # prevenir OOM en GPU
+
+    # Construye los mensajes para la generación
+    messages = [
+        {"role": "system", "content": SYS_PROMPT},
+        {"role": "user", "content": formatted_prompt}
+    ]
+    
+    # Prepara la entrada usando la plantilla de chat del tokenizer
     input_ids = tokenizer.apply_chat_template(
-      messages,
-      add_generation_prompt=True,
-      return_tensors="pt"
+        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
-        )
+    
+    # Define el streamer para obtener la salida progresivamente
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        input_ids= input_ids,
+        input_ids=input_ids,
         streamer=streamer,
         max_new_tokens=1024,
         do_sample=True,
@@ -95,46 +115,35 @@ def talk(prompt,history=[]):
         temperature=0.75,
         eos_token_id=terminators,
     )
+    
+    # Ejecuta la generación en un hilo para ir transmitiendo resultados
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
-
+    
     outputs = []
     for text in streamer:
         outputs.append(text)
         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 
-"""
-
-
-demo = gr.ChatInterface(
-    fn=talk,
-    chatbot=gr.Chatbot(
-        show_label=True,
-        show_share_button=True,
-        show_copy_button=True,
-        layout="bubble",
-        bubble_full_width=False,
-    ),
-    theme="Soft",
-    type="tuples",
-    examples=["what's anarchy ? "],
-    title=TITLE,
-    description=DESCRIPTION,
-    autofocus=False,
-    autoscroll = False,
+# Interfaz de Gradio: se añade un dropdown para la selección del modelo
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# RAG Chatbot")
+    gr.Markdown("Selecciona el modelo de inteligencia artificial:")
     
-)
-demo.launch(debug=True)
+    with gr.Row():
+        # Dropdown para elegir el modelo
+        modelo_selector = gr.Dropdown(choices=list(MODELOS.keys()), value="Llama-3.2-3B-Instruct", label="Modelo")
+    
+    chatbot = gr.Chatbot(show_label=True, show_share_button=True, show_copy_button=True, layout="bubble")
+    
+    # Caja de entrada para la consulta del usuario
+    prompt_input = gr.Textbox(lines=2, label="Ingresa tu pregunta")
+    
+    # Botón para enviar la consulta
+    send_btn = gr.Button("Enviar")
+    
+    # Llamada a la función talk, pasando el prompt y la selección de modelo
+    send_btn.click(fn=talk, inputs=[prompt_input, modelo_selector], outputs=chatbot, api_name="chatbot")
+    
+if __name__ == "__main__":
+    demo.launch(debug=True)