app.py

import streamlit as st
import os
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceHubEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import HuggingFaceHub
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate

st.set_page_config(page_title='preguntaDOC')
st.header("Pregunta a tu PDF")

# Campo para el token de Hugging Face
huggingface_api_token = st.text_input('Hugging Face API Token (requerido)', type='password')

# Selección de modelo
modelo_options = {
    "Google FLAN-T5 Small": "google/flan-t5-small",
    "Google FLAN-T5 Base": "google/flan-t5-base",
    "BLOOM 560M": "bigscience/bloom-560m",
    "BLOOM 1.1B": "bigscience/bloom-1b1",
    "Falcon 7B Instruct": "tiiuae/falcon-7b-instruct",
    "Gemma 2B": "google/gemma-2b",
    "Gemma 2B Instruct": "google/gemma-2b-it"
}

selected_model = st.selectbox("Selecciona un modelo:", list(modelo_options.keys()))
modelo_id = modelo_options[selected_model]

pdf_obj = st.file_uploader("Carga tu documento", type="pdf", on_change=st.cache_resource.clear)

@st.cache_resource 
def create_embeddings(pdf, api_token):
    if not api_token:
        st.error("Se requiere un token de API de Hugging Face")
        return None
        
    os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_token
    
    pdf_reader = PdfReader(pdf)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=100,
        length_function=len
    )        
    chunks = text_splitter.split_text(text)
    
    # Usar HuggingFaceHubEmbeddings
    embeddings = HuggingFaceHubEmbeddings(
        repo_id="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        huggingfacehub_api_token=api_token
    )
    
    knowledge_base = FAISS.from_texts(chunks, embeddings)
    return knowledge_base

if pdf_obj and huggingface_api_token:
    knowledge_base = create_embeddings(pdf_obj, huggingface_api_token)
    
    if knowledge_base:
        user_question = st.text_input("Haz una pregunta sobre tu PDF:")
        
        if user_question:
            docs = knowledge_base.similarity_search(user_question, 3)
            
            # Configurar los parámetros del modelo según el tipo
            model_kwargs = {}
            
            # Verificar el tipo de modelo para usar los parámetros adecuados
            if "flan-t5" in modelo_id:
                model_kwargs = {"temperature": 0.5, "max_length": 512}
            elif "bloom" in modelo_id:
                model_kwargs = {"temperature": 0.7, "max_length": 512}
            elif "falcon" in modelo_id or "llama" in modelo_id or "gemma" in modelo_id:
                model_kwargs = {"temperature": 0.1, "max_new_tokens": 512}
            else:
                model_kwargs = {"temperature": 0.5, "max_length": 512}
            
            # Crear el LLM con los parámetros adecuados
            llm = HuggingFaceHub(
                repo_id=modelo_id,
                huggingfacehub_api_token=huggingface_api_token,
                model_kwargs=model_kwargs
            )
            
            # Prompt diferente según el tipo de modelo
            if "falcon" in modelo_id or "llama" in modelo_id or "gemma" in modelo_id:
                prompt_template = """
                <|system|>
                Responde a la siguiente pregunta basándote únicamente en el contexto proporcionado.
                </|system|>
                
                <|user|>
                Contexto: {context}
                
                Pregunta: {question}
                </|user|>
                
                <|assistant|>
                """
            else:
                prompt_template = """
                Responde a la siguiente pregunta basándote únicamente en el contexto proporcionado.
                
                Contexto: {context}
                
                Pregunta: {question}
                
                Respuesta:
                """
            
            PROMPT = PromptTemplate(
                template=prompt_template, 
                input_variables=["context", "question"]
            )
            
            chain = load_qa_chain(llm, chain_type="stuff", prompt=PROMPT)
            
            with st.spinner(f"Procesando tu pregunta con {selected_model}..."):
                try:
                    respuesta = chain.run(input_documents=docs, question=user_question)
                    st.write(respuesta)
                except Exception as e:
                    st.error(f"Error al procesar tu pregunta: {str(e)}")
                    st.info("Sugerencia: Intenta con un modelo diferente. Algunos modelos pueden requerir más recursos o tener limitaciones específicas.")
elif not huggingface_api_token and pdf_obj:
    st.warning("Por favor, ingresa tu token de API de Hugging Face para continuar.")