app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from huggingface_hub import login
from PyPDF2 import PdfReader
from docx import Document
import csv
import json
import os
import torch

huggingface_token = os.getenv('HUGGINGFACE_TOKEN')

# Realizar el inicio de sesión de Hugging Face solo si el token está disponible
if huggingface_token:
    login(token=huggingface_token)

# Configuración del modelo
@st.cache_resource
def load_llm():
    llm = HuggingFaceEndpoint(
        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
        task="text-generation"
    )
    llm_engine_hf = ChatHuggingFace(llm=llm)
    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
    return llm_engine_hf, tokenizer

llm_engine_hf, tokenizer = load_llm()

# Configuración del modelo de clasificación
@st.cache_resource
def load_classification_model():
    tokenizer = AutoTokenizer.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
    model = AutoModelForSequenceClassification.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
    return model, tokenizer

classification_model, classification_tokenizer = load_classification_model()

id2label = {0: "multas", 1: "politicas_de_privacidad", 2: "contratos", 3: "denuncias", 4: "otros"}

def classify_text(text):
    inputs = classification_tokenizer(text, return_tensors="pt", max_length=4096, truncation=True, padding="max_length")
    classification_model.eval()
    with torch.no_grad():
        outputs = classification_model(**inputs)
    logits = outputs.logits
    predicted_class_id = logits.argmax(dim=-1).item()
    predicted_label = id2label[predicted_class_id]
    return f"Clasificación: {predicted_label}\n\nDocumento:\n{text}"

def translate(text, target_language):
    template = '''
    Por favor, traduzca el siguiente documento al {LANGUAGE}:
<document>
{TEXT}
</document>
Asegúrese de que la traducción sea precisa y conserve el significado original del documento.
    '''
    
    formatted_prompt = template.replace("{TEXT}", text).replace("{LANGUAGE}", target_language)
    response = llm_engine_hf.invoke(formatted_prompt)
    translated_text = response.content
    
    return translated_text

def summarize(text, length):
    template = f'''
    Por favor, haga un resumen {length} del siguiente documento:
<document>
{text}
</document>
Asegúrese de que el resumen sea conciso y conserve el significado original del documento.
    '''
    
    response = llm_engine_hf.invoke(template)
    summarized_text = response.content
    
    return summarized_text

def handle_uploaded_file(uploaded_file):
    try:
        if uploaded_file.name.endswith(".txt"):
            text = uploaded_file.read().decode("utf-8")
        elif uploaded_file.name.endswith(".pdf"):
            reader = PdfReader(uploaded_file)
            text = ""
            for page in range(len(reader.pages)):
                text += reader.pages[page].extract_text()
        elif uploaded_file.name.endswith(".docx"):
            doc = Document(uploaded_file)
            text = "\n".join([para.text for para in doc.paragraphs])
        elif uploaded_file.name.endswith(".csv"):
            text = ""
            content = uploaded_file.read().decode("utf-8").splitlines()
            reader = csv.reader(content)
            text = " ".join([" ".join(row) for row in reader])
        elif uploaded_file.name.endswith(".json"):
            data = json.load(uploaded_file)
            text = json.dumps(data, indent=4)
        else:
            text = "Tipo de archivo no soportado."
        return text
    except Exception as e:
        return str(e)

def main():
    st.set_page_config(page_title="LexAIcon", page_icon="utils/icon.jpg", layout="centered", initial_sidebar_state="auto", menu_items=None)

    st.image("./icon.jpg", width=100)

    st.title("LexAIcon: Traduce, resume y explica textos legales")
    st.write("Puedes conversar con este chatbot basado en Mistral7B-Instruct y subir archivos para que el chatbot los procese.")

    if "generated" not in st.session_state:
        st.session_state["generated"] = []

    # Entrada del usuario
    user_input = st.text_input("Tú: ", "")

    # Botones de Resumir, Traducir y Explicar
    operation = st.radio("Selecciona una operación", ["Resumir", "Traducir", "Explicar"])

    target_language = None
    summary_length = None

    if operation == "Traducir":
        target_language = st.selectbox("Selecciona el idioma de traducción", ["español", "inglés", "francés", "alemán"])

    if operation == "Resumir":
        summary_length = st.selectbox("Selecciona la longitud del resumen", ["corto", "medio", "largo"])

    # Manejo de archivos subidos
    uploaded_files = st.file_uploader("Sube un archivo", type=["txt", "pdf", "docx", "csv", "json"], accept_multiple_files=True)

    if st.button("Enviar"):
        if user_input:
            response = llm_engine_hf.invoke(user_input)
            st.session_state.generated.append({"user": user_input, "bot": response['generated_text']})

    if st.button("Ejecutar"):
        if uploaded_files:
            for uploaded_file in uploaded_files:
                file_content = handle_uploaded_file(uploaded_file)
                if operation == "Resumir":
                    if summary_length == "corto":
                        length = "de aproximadamente 50 palabras"
                    elif summary_length == "medio":
                        length = "de aproximadamente 100 palabras"
                    elif summary_length == "largo":
                        length = "de aproximadamente 500 palabras"
                    result = summarize(file_content, length)
                elif operation == "Traducir":
                    result = translate(file_content, target_language)
                elif operation == "Explicar":
                    result = classify_text(file_content)
                st.write(result)

    if st.session_state.get("generated"):
        for chat in st.session_state["generated"]:
            st.write(f"Tú: {chat['user']}")
            st.write(f"Chatbot: {chat['bot']}")

if __name__ == "__main__":
    main()