Update app.py

app.py

@@ -1,106 +1,53 @@
 import os
 import requests
-import re
-from PyPDF2 import PdfReader
-from sentence_transformers import SentenceTransformer, util
+from llama_index.core import VectorStoreIndex, Settings
+from llama_index.readers.file import PDFReader
 import gradio as gr
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+# Disable the default LLM
+Settings.llm = None
 
-# 1. Descargar el PDF
 def download_pdf(url, destination):
-    """Descarga un PDF desde una URL y lo guarda en la ruta especificada."""
     os.makedirs(os.path.dirname(destination), exist_ok=True)
     response = requests.get(url)
     with open(destination, 'wb') as f:
         f.write(response.content)
 
-# 2. Extraer los artículos del PDF
-def extract_articles_from_pdf(pdf_path):
-    """Extrae artículos del PDF basado en el formato del Código Penal."""
-    reader = PdfReader(pdf_path)
-    text = ""
-    for page in reader.pages:
-        text += page.extract_text()
-
-    # Usar regex para segmentar los artículos
-    article_pattern = r'(Artículo \d+\..*?)(?=Artículo \d+\.|$)'
-    matches = re.findall(article_pattern, text, re.DOTALL)
-
-    # Crear un diccionario de artículos
-    articles = {}
-    for match in matches:
-        lines = match.strip().split("\n")
-        title = lines[0].strip()  # Ejemplo: "Artículo 138."
-        content = " ".join(line.strip() for line in lines[1:]).strip()
-        articles[title] = content
-
-    return articles
-
-# 3. Crear embeddings para los artículos
-def create_article_embeddings(articles, model_name="paraphrase-multilingual-mpnet-base-v2"):
-    """Crea embeddings para los artículos utilizando SentenceTransformers."""
-    model = SentenceTransformer(model_name)
-    article_keys = list(articles.keys())
-    article_embeddings = model.encode(list(articles.values()), convert_to_tensor=True)
-    return article_keys, article_embeddings, model
-
-# 4. Buscar el artículo relevante
-def find_article(question, article_keys, article_embeddings, model, articles):
-    # Filtrar artículos relevantes usando palabras clave
-    keywords = question.lower().split()  # Dividir pregunta en palabras clave
-    filtered_articles = {
-        key: value for key, value in articles.items()
-        if any(keyword in value.lower() for keyword in keywords)
-    }
-
-    if not filtered_articles:
-        # Si no hay artículos relevantes basados en palabras clave, usar todos
-        filtered_articles = articles
-
-    # Crear nuevos embeddings para los artículos filtrados
-    filtered_keys = list(filtered_articles.keys())
-    filtered_embeddings = model.encode(list(filtered_articles.values()), convert_to_tensor=True)
-
-    # Calcular similitud con la pregunta
-    question_embedding = model.encode(question, convert_to_tensor=True)
-    scores = util.pytorch_cos_sim(question_embedding, filtered_embeddings)
-    best_match_idx = scores.argmax()
-    best_article_key = filtered_keys[best_match_idx]
-    return f"{best_article_key}\n{filtered_articles[best_article_key]}"
-
-
-# Flujo principal
-def main():
-    # Configuración inicial
-    pdf_url = 'https://www.boe.es/buscar/pdf/1995/BOE-A-1995-25444-consolidado.pdf'
-    pdf_path = './BOE-A-1995-25444-consolidado.pdf'
-
-    # Descargar el PDF si no existe
-    if not os.path.exists(pdf_path):
-        print("Descargando el Código Penal...")
-        download_pdf(pdf_url, pdf_path)
-
-    # Extraer y procesar los artículos
-    print("Extrayendo artículos del Código Penal...")
-    articles = extract_articles_from_pdf(pdf_path)
-
-    # Crear embeddings para los artículos
-    print("Creando embeddings para los artículos...")
-    article_keys, article_embeddings, model = create_article_embeddings(articles)
-
-    # Función para responder preguntas
-    def search_law(query):
-        return find_article(query, article_keys, article_embeddings, model, articles)
-
-    # Iniciar la interfaz de Gradio
-    print("Lanzando la aplicación...")
-    gr.Interface(
-        fn=search_law, 
-        inputs="text", 
-        outputs="text",
-        title="Búsqueda en el Código Penal Español",
-        description="Realiza preguntas sobre delitos y penas en el Código Penal Español."
-    ).launch()
+def create_index_from_pdf(pdf_path):
+    pdf_reader = PDFReader()
+    documents = pdf_reader.load_data(file=pdf_path)
     
-
-if __name__ == "__main__":
-    main()
+    embed_model = HuggingFaceEmbedding(model_name='BAAI/bge-large-es')
+    
+    index = VectorStoreIndex.from_documents(
+        documents,
+        embed_model=embed_model
+    )
+    
+    query_engine = index.as_query_engine(
+        similarity_top_k=3,  # Increased to find more relevant context
+        response_mode="compact"
+    )
+    
+    return query_engine
+
+pdf_url = 'https://www.boe.es/buscar/pdf/1995/BOE-A-1995-25444-consolidado.pdf'
+pdf_path = './BOE-A-1995-25444-consolidado.pdf'
+
+download_pdf(pdf_url, pdf_path)
+query_engine = create_index_from_pdf(pdf_path)
+
+def search_pdf(query):
+    # Modificar la consulta para buscar específicamente penas
+    modified_query = f"Pena para el delito de {query}"
+    response = query_engine.query(modified_query)
+    return response.response
+
+gr.Interface(
+    fn=search_pdf, 
+    inputs="text", 
+    outputs="text",
+    title="Buscador de Penas en Código Penal",
+    description="Introduce un tipo de delito para encontrar su pena correspondiente"
+).launch()