Update app.py

app.py

@@ -1,18 +1,13 @@
+import os
 import gradio as gr
-from langchain_community.llms import HuggingFaceHub
-from langchain_community.vectorstores import Chroma
-from langchain_community.document_loaders import PyPDFLoader
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_core.output_parsers import StrOutputParser
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain import hub
+from langchain_community.vectorstores import Chroma
 from rerankers import Reranker
-import os
 
-# Configuración del token de acceso a Hugging Face (si usas modelo privado)
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
-# Cargar PDF
+# Cargar PDF y partirlo en chunks
 loader = PyPDFLoader("80dias.pdf")
 documents = loader.load()
 splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=20)
@@ -23,44 +18,54 @@ embedding_model = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
 vectordb = Chroma.from_documents(splits, embedding=embeddings)
 
-# Modelo LLM desde HuggingFace (usa uno disponible en Spaces)
-llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.1", model_kwargs={"temperature": 0.5, "max_new_tokens": 500})
-chain = llm | StrOutputParser()
-
-# Reranker
+# Inicializar reranker
 ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type="colbert")
 
-# Función RAG
+# Cargar modelo de lenguaje de Hugging Face
+model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# Función principal RAG
 def rag_chat(message, history):
-    # Solo usamos el mensaje del usuario
     query = message
-
     results = vectordb.similarity_search_with_score(query)
+    
+    # Seleccionar contextos relevantes
     context = []
     for doc, score in results:
         if score < 7:
             context.append(doc.page_content)
+
     if not context:
         return "No tengo información suficiente para responder a esa pregunta."
 
+    # Aplicar reranking
     ranking = ranker.rank(query=query, docs=context)
     best_context = ranking[0].text
 
-    prompt = f"""Contesta a la siguiente pregunta usando solo el contexto que se proporciona:
+    # Crear prompt final
+    prompt = f"""Responde a la siguiente pregunta utilizando solo el contexto proporcionado:
 
-    Contexto:
-    {best_context}
+Contexto:
+{best_context}
 
-    Pregunta: {query}
-    Respuesta:"""
+Pregunta: {query}
+Respuesta:"""
 
-    return llm.invoke(prompt)
-    
-# Interfaz Gradio
+    # Generar respuesta
+    output = generator(prompt, max_new_tokens=300, do_sample=False)[0]["generated_text"]
+    response = output.split("Respuesta:")[-1].strip()
+    return response
+
+# Gradio Chat Interface
 iface = gr.ChatInterface(
     fn=rag_chat,
     title="Chat Julio Verne - RAG",
     description="Pregunta lo que quieras sobre *La vuelta al mundo en 80 días* de Julio Verne.",
-    chatbot=gr.Chatbot(type="messages")  # 👈 Esto elimina el warning de formato obsoleto
+    chatbot=gr.Chatbot(type="messages"),
+    theme="default"
 )
+
 iface.launch()