app.py

import PyPDF2 as pypdf2
import os
import re
from pathlib import Path
from sentence_transformers import SentenceTransformer
import requests
from chromadb import Client
from transformers import AutoModelForQuestionAnswering, AutoTokenizer
import gradio as gr
import torch



# ==================== CONFIGURACIÓN ====================

DATA_DIR = "./"
CHUNK_SIZE = 500
OVERLAPPING = 100
N_SEARCH_RESULTS = 5

# ==================== FUNCIONES AUXILIARES ====================

def extraer_texto_pdf(pdf_path):
    """Extraer texto de un archivo PDF"""
    try:
        with open(pdf_path, "rb") as file:
            reader = pypdf2.PdfReader(file)
            text = ""
            for page in reader.pages:
                text += page.extract_text()
        return text
    except Exception as e:
        print(f"Error extrayendo texto de {pdf_path}: {str(e)}")
        return ""

def limpiar_texto(texto):
    """
    Limpia el texto de espacios múltiples, saltos de línea y tabs
    """
    # Remover tabs
    texto = re.sub(r'\t', ' ', texto)
    
    # Remover saltos de línea
    texto = re.sub(r'\n+', ' ', texto)
    
    # Remover espacios múltiples (3 o más)
    texto = re.sub(r' {3,}', ' ', texto)
    
    # Remover espacios dobles
    texto = re.sub(r'  +', ' ', texto)
    
    # Remover espacios al inicio y final
    texto = texto.strip()
    
    # Remover espacios antes de puntos
    texto = re.sub(r' \.', '.', texto)
    
    # Remover puntos repetidos
    texto = re.sub(r'\.{2,}', '.', texto)
    
    return texto

def dividir_texto(text, chunk_size=CHUNK_SIZE, overlapping=OVERLAPPING):
    """
    Dividir texto en fragmentos con superposición
    """
    chunks = []
    start = 0
    
    while start < len(text):
        end = start + chunk_size
        
        if end < len(text):
            # Buscar el último espacio para no cortar palabras
            last_space = text.rfind(' ', start, end)
            if last_space > start:
                end = last_space
        
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        
        start = end - overlapping
        
        if start >= len(text):
            break
    
    return chunks

def cargar_y_procesar_pdfs():
    """Cargar, procesar y crear embeddings de los PDFs"""
    print("\n Leyendo los pdfs del reglamento")
    
    # Leer PDFs
    print("\nLeyendo archivos PDF...")
    pdf_folder = DATA_DIR
    pdf_files = [f for f in os.listdir(pdf_folder) if f.endswith('.pdf')]
    
    if not pdf_files:
        return None, None, None, None
    
    pdf_texts = []
    for pdf_file in pdf_files:
        text = extraer_texto_pdf(os.path.join(pdf_folder, pdf_file))
        if text:
            pdf_texts.append(text)
    
    # Limpiar textos
    print("Limpiando textos")
    pdf_texts_limpios = [limpiar_texto(text) for text in pdf_texts]
    
    # Dividir en chunks
    print("Dividiendo textos en fragmentos")
    chunks = []
    for texto in pdf_texts_limpios:
        chunks.extend(dividir_texto(texto, chunk_size=CHUNK_SIZE, overlapping=OVERLAPPING))
    
    print(f"✓ Total de fragmentos generados: {len(chunks)}")
    
    # Crear embeddings
    print("Generando embeddings")
    embedding_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
    embeddings = embedding_model.encode(chunks)
    print(f"✓ Dimensión de embeddings: {embeddings.shape}")
    
    # Crear base de datos vectorial
    print("Creando base de datos vectorial")
    client = Client()
    collection = client.create_collection(
        name="documentos_pdf",
        metadata={"hnsw:space": "cosine"}
    )
    
    # Agregar fragmentos a la colección
    collection.add(
        embeddings=embeddings,
        metadatas=[{"source": "pdf", "chunk_id": i} for i in range(len(chunks))],
        documents=chunks,
        ids=[f"chunk_{i}" for i in range(len(chunks))]
    )
    
    print(f"Base de datos vectorial creada con {len(chunks)} fragmentos\n")
    
    return embedding_model, collection, client, chunks

# ==================== INICIALIZACIÓN ====================

print("Cargando los documentos y los modelos necesarios")

embedding_model, collection, client, chunks = cargar_y_procesar_pdfs()

# Cargar modelo QA
print("Cargando modelo de Preguntas y Respuestas")
model_name = 'mrm8488/distill-bert-base-spanish-wwm-cased-finetuned-spa-squad2-es'
tokenizer = AutoTokenizer.from_pretrained(model_name)
qa_model = AutoModelForQuestionAnswering.from_pretrained(model_name)
print("Modelo QA cargado\n")

# ==================== FUNCIÓN PRINCIPAL ====================

def obtener_respuesta(pregunta):
    if not pregunta.strip():
        return "La pregunta no es valida."
    
    if embedding_model is None or collection is None:
        return "Los documentos no se cargaron correctamente."
    
    try:
        # Generar embedding de la pregunta
        query_embedding = embedding_model.encode(pregunta)
        
        # Buscar documentos similares
        res = collection.query(
            query_embeddings=[query_embedding],
            n_results=N_SEARCH_RESULTS
        )
        
        # Combinar contexto
        context = " ".join(res['documents'][0])
        
        # Tokenizar pregunta y contexto
        inputs = tokenizer(pregunta, context, return_tensors='pt', max_length=512, truncation=True)
        
        # Obtener predicciones del modelo
        with torch.no_grad():
            outputs = qa_model(**inputs)
        
        # Extraer índices de inicio y fin
        start_idx = torch.argmax(outputs.start_logits)
        end_idx = torch.argmax(outputs.end_logits) + 1
        
        # Decodificar respuesta
        answer_tokens = inputs.input_ids[0, start_idx:end_idx]
        answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
        
        return answer if answer.strip() else "No se encontró una respuesta clara."
    
    except Exception as e:
        return f"Error procesando la pregunta: {str(e)}"

# ==================== INTERFAZ GRADIO ====================

with gr.Blocks(title="Asistente de Reglamentos UNAB") as demo:
    gr.Markdown("""
    # Asistente de Reglamentos UNAB
    
    Realiza preguntas sobre los reglamentos de la Universidad Autónoma de Bucaramanga
    utilizando búsqueda semántica y procesamiento de lenguaje natural.
    """)
    
    with gr.Row():
        with gr.Column():
            question_input = gr.Textbox(
                lines=3,
                placeholder="Ejemplo: ¿Cuál es el proceso de inscripción?",
                label="Tu pregunta"
            )
            submit_button = gr.Button("Buscar respuesta", variant="primary")
        
        with gr.Column():
            answer_output = gr.Textbox(
                label="Respuesta",
                lines=5,
                interactive=False
            )
    
    submit_button.click(
        fn=obtener_respuesta,
        inputs=question_input,
        outputs=answer_output
    )
    
    # Enter key support
    question_input.submit(
        fn=obtener_respuesta,
        inputs=question_input,
        outputs=answer_output
    )
    
    gr.Markdown("""
    ### Información
    - Este asistente responde preguntas basadas en documentos de la UNAB
    - Utiliza búsqueda semántica para encontrar fragmentos relevantes
    - Las respuestas se generan automáticamente mediante IA
    - Para mejores resultados, sé específico en tus preguntas
    """)

if __name__ == "__main__":
    demo.launch(share=False)