Update src/model_load.py

src/model_load.py

@@ -1,61 +1,60 @@
-from langchain.chains import RetrievalQAWithSourcesChain
-from langchain.llms import HuggingFacePipeline
-from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM, BitsAndBytesConfig
-import torch
-from langchain.prompts import PromptTemplate
-from langchain.llms import HuggingFaceHub
-from langchain.chains import LLMChain
-
-def load_model():
-    model_name="tiiuae/Falcon3-10B-Instruct"
-
-    # Cargar tokenizer y modelo de Hugging Face
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name,
-                                                 torch_dtype=torch.float16,
-                                                 device_map="auto",
-                                                 quantization_config=BitsAndBytesConfig(load_in_8bit=True))
-    
-    # Crear pipeline de generación de texto
-    text_generation_pipeline = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=128,
-    repetition_penalty=1.2,
-    device_map="auto"
-    )
-    # Crear el LLM compatible con LangChain
-    llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
-
-    # Crear la plantilla de prompt que tomará el texto y la pregunta
-    prompt_template = """
-    Dado el siguiente texto extraído de varios documentos y una pregunta, crea una respuesta utilizando la información proporcionada. Si la pregunta sale por fuera de la información proporcionada responde con "No tengo información al respecto" y corta la respuesta.
-
-    **Documentos relevantes:**
-    {documento}
-
-    **Pregunta:**
-    {pregunta}
-
-    **Respuesta:**
-    """
-
-    # Crear el prompt con las variables necesarias
-    prompt = PromptTemplate(input_variables=["documento", "pregunta"], template=prompt_template)
-
-    # Crear una cadena de LLMChain que combine el retriever y el prompt
-    qa_chain = LLMChain(llm=llm, prompt=prompt)
-
-def ask(pregunta: str,retriever):
-
-    #Busqueda de documentos mediante el retriever
-    documentos=retriever.get_relevant_documents(pregunta)
-
-    #Generacion de la respuesta
-    respuesta = qa_chain.invoke({
-    "documentos": "\n".join([doc.page_content for doc in docs_relevantes]),
-    "pregunta": pregunta
-    })
-
+from langchain.chains import RetrievalQAWithSourcesChain
+from langchain.llms import HuggingFacePipeline
+from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+from langchain.prompts import PromptTemplate
+from langchain.llms import HuggingFaceHub
+from langchain.chains import LLMChain
+
+def load_model():
+    model_name="tiiuae/Falcon3-10B-Instruct"
+
+    # Cargar tokenizer y modelo de Hugging Face
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name,
+                                                 torch_dtype=torch.float16,
+                                                 device_map="auto")
+    
+    # Crear pipeline de generación de texto
+    text_generation_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=128,
+    repetition_penalty=1.2,
+    device_map="auto"
+    )
+    # Crear el LLM compatible con LangChain
+    llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
+
+    # Crear la plantilla de prompt que tomará el texto y la pregunta
+    prompt_template = """
+    Dado el siguiente texto extraído de varios documentos y una pregunta, crea una respuesta utilizando la información proporcionada. Si la pregunta sale por fuera de la información proporcionada responde con "No tengo información al respecto" y corta la respuesta.
+
+    **Documentos relevantes:**
+    {documento}
+
+    **Pregunta:**
+    {pregunta}
+
+    **Respuesta:**
+    """
+
+    # Crear el prompt con las variables necesarias
+    prompt = PromptTemplate(input_variables=["documento", "pregunta"], template=prompt_template)
+
+    # Crear una cadena de LLMChain que combine el retriever y el prompt
+    qa_chain = LLMChain(llm=llm, prompt=prompt)
+
+def ask(pregunta: str,retriever):
+
+    #Busqueda de documentos mediante el retriever
+    documentos=retriever.get_relevant_documents(pregunta)
+
+    #Generacion de la respuesta
+    respuesta = qa_chain.invoke({
+    "documentos": "\n".join([doc.page_content for doc in docs_relevantes]),
+    "pregunta": pregunta
+    })
+
     return respuesta["text"].split("**Respuesta:**\n")[1]