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-7B-Instruct"
    max_memory = {0: "24GB", "cpu": "30GB"}
    # Cargar tokenizer y modelo de Hugging Face
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name,
                                                 torch_dtype=torch.bfloat16).to("cuda")
    
    # 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 = prompt | llm
    return qa_chain
def ask(pregunta: str,retriever,qa_chain):

    
    
    #Busqueda de documentos mediante el retriever
    documentos=retriever.invoke(pregunta)

    #Generacion de la respuesta
    respuesta = qa_chain.invoke({
    "documento": "\n".join([doc.page_content for doc in documentos]),
    "pregunta": pregunta
    })

    return respuesta