app.py

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain_community.llms import HuggingFacePipeline
from langchain_core.output_parsers import StrOutputParser
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain import hub
import gradio as gr

# ------------------------------
# MODELO
# ------------------------------
model_id = "mistralai/Mistral-7B-Instruct-v0.1"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.7
)

llm = HuggingFacePipeline(pipeline=pipe)
parser = StrOutputParser()

# ------------------------------
# EMBEDDINGS + CHROMA
# ------------------------------
embedding_function = HuggingFaceEmbeddings(
    model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
    model_kwargs={"device": "cpu"}
)

vectordb = Chroma(
    persist_directory="chroma_db",
    embedding_function=embedding_function
)

# ------------------------------
# FUNCIÓN RAG
# ------------------------------
def responder_pregunta(query):
    docs = vectordb.similarity_search_with_score(query, k=5)
    prompt = hub.pull("rlm/rag-prompt")
    rag_chain = prompt | llm | parser

    context = []
    for doc, score in docs:
        if score < 7:
            context.append(doc.page_content)

    if context:
        context_text = "\n".join(context)
        result = rag_chain.invoke({"context": context_text, "question": query})
        return result
    else:
        return "No tengo información suficiente para responder a esta pregunta."

# ------------------------------
# INTERFAZ GRADIO
# ------------------------------
gr.Interface(
    fn=responder_pregunta,
    inputs=gr.Textbox(label="Pregunta sobre nutrición"),
    outputs="text",
    title="Sistema RAG sobre Nutrición Clínica",
    description="Haz preguntas sobre el manual clínico procesado con embeddings + Mistral 7B."
).launch()