Update app.py

app.py

@@ -1,69 +1,168 @@
-from huggingface_hub import InferenceClient
-import gradio as gr
-import json
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
 import PyPDF2
+import gradio as gr
+from langchain.prompts import PromptTemplate
+from langchain.chains.summarize import load_summarize_chain
+from pathlib import Path
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.3"
+# Configuración del modelo de resumen
+llm = HuggingFaceEndpoint(
+    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
+    task="text-generation",
+    max_new_tokens=4096,
+    temperature=0.5,
+    do_sample=False,
 )
+llm_engine_hf = ChatHuggingFace(llm=llm)
+
+# Configuración del modelo de clasificación
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
+model = AutoModelForSequenceClassification.from_pretrained("mrm8488/legal-longformer-base-8192-spanish")
+
+id2label = {0: "multas", 1: "politicas_de_privacidad", 2: "contratos", 3: "denuncias", 4: "otros"}
+
+def read_pdf(file_path):
+    pdf_reader = PyPDF2.PdfReader(file_path)
+    text = ""
+    for page in range(len(pdf_reader.pages)):
+        text += pdf_reader.pages[page].extract_text()
+    return text
 
-def format_prompt(mode, message, instructions, history):
-    prompt = f"<s>[MODE] {mode} [/MODE] "
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST] {bot_response}</s> "
-    prompt += f"[INST] {message} {instructions} [/INST]"
-    return prompt
-
-def process_input(file, file_type):
-    if file_type == 'pdf':
-        reader = PyPDF2.PdfFileReader(file.name)
-        text = ""
-        for page in range(reader.numPages):
-            text += reader.getPage(page).extract_text()
-        return text
-    elif file_type == 'json':
-        with open(file.name, 'r') as f:
-            data = json.load(f)
-        return json.dumps(data, indent=2)
-    return file
-
-def generate(mode, file, file_type, instructions, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
-    temperature = max(float(temperature), 1e-2)
-    top_p = float(top_p)
+def summarize(file):
+    # Leer el contenido del archivo subido
+    file_path = file.name
+    if file_path.endswith('.pdf'):
+        text = read_pdf(file_path)
+    else:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            text = f.read()
+            
+    template = '''
+    Please carefully read the following document:
+<document>
+{TEXT}
+</document>
+After reading through the document, identify the key points and main ideas covered in the text. Organize these key points into a concise bulleted list that summarizes the essential information from the document. The summary should have a maximum of 10 bullet points.
+Your goal is to be comprehensive in capturing the core content of the document, while also being concise in how you express each summary point. Omit minor details and focus on the central themes and important facts.
+    '''
     
-    text_input = process_input(file, file_type)
-    formatted_prompt = format_prompt(mode, text_input, instructions, history)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
+    prompt = PromptTemplate(
+        template=template,
+        input_variables=['TEXT']    
     )
+    
+    formatted_prompt = prompt.format(TEXT=text)
+    output_summary = llm_engine_hf.invoke(formatted_prompt)
+    
+    return output_summary.content
 
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
+def classify_text(text):
+    inputs = tokenizer(text, return_tensors="pt", max_length=4096, truncation=True, padding="max_length")
+    model.eval()
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_id = logits.argmax(dim=-1).item()
+    predicted_label = id2label[predicted_class_id]
+    return predicted_label
 
-    for response in stream:
-        output += response.token.text
-        yield output
-    return output
+def translate(file, target_language):
+    # Leer el contenido del archivo subido
+    file_path = file.name
+    if file_path.endswith('.pdf'):
+        text = read_pdf(file_path)
+    else:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            text = f.read()
 
+    template = '''
+    Please translate the following document to {LANGUAGE}:
+<document>
+{TEXT}
+</document>
+Ensure that the translation is accurate and preserves the original meaning of the document.
+    '''
+    
+    prompt = PromptTemplate(
+        template=template,
+        input_variables=['TEXT', 'LANGUAGE']    
+    )
+    
+    formatted_prompt = prompt.format(TEXT=text, LANGUAGE=target_language)
+    translated_text = llm_engine_hf.invoke(formatted_prompt)
+    
+    return translated_text.content
+
+def process_file(file, action, target_language=None):
+    if action == "Resumen":
+        return summarize(file)
+    elif action == "Clasificar":
+        file_path = file.name
+        if file_path.endswith('.pdf'):
+            text = read_pdf(file_path)
+        else:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                text = f.read()
+        return classify_text(text)
+    elif action == "Traducir":
+        return translate(file, target_language)
+    else:
+        return "Acción no válida"
+
+def download_text(output_text, filename='output.txt'):
+    if output_text:
+        file_path = Path(filename)
+        with open(file_path, 'w', encoding='utf-8') as f:
+            f.write(output_text)
+        return file_path
+    else:
+        return None
+
+def create_download_file(output_text, filename='output.txt'):
+    file_path = download_text(output_text, filename)
+    return str(file_path) if file_path else None
+
+# Crear la interfaz de Gradio
 with gr.Blocks() as demo:
-    mode = gr.Dropdown(label="Mode", choices=["translation", "summary", "explanation"], value="translation")
-    file = gr.File(label="Input File", type="file")
-    file_type = gr.Radio(label="File Type", choices=["pdf", "json"], value="pdf")
-    instructions = gr.Textbox(label="Additional Instructions", placeholder="Enter any additional instructions here")
-    chatbot = gr.Chatbot()
-
-    def update_chatbot(mode, file, file_type, instructions, history):
-        return generate(mode, file, file_type, instructions, history)
+    gr.Markdown("## Document Processor")
+
+    with gr.Row():
+        with gr.Column():
+            file = gr.File(label="Subir un archivo")
+            action = gr.Radio(label="Seleccione una acción", choices=["Resumen", "Clasificar", "Traducir"])
+            target_language = gr.Dropdown(label="Seleccionar idioma de traducción", choices=["en", "fr", "de"], visible=False)
+        
+        with gr.Column():
+            output_text = gr.Textbox(label="Resultado", lines=20)
+
+    def update_language_dropdown(action):
+        if action == "Traducir":
+            return gr.update(visible=True)
+        else:
+            return gr.update(visible=False)
     
-    gr.Interface(
-        fn=update_chatbot,
-        inputs=[mode, file, file_type, instructions, chatbot],
-        outputs=chatbot,
-        title="Mistral 7B v0.3"
-    ).launch()
+    action.change(update_language_dropdown, inputs=action, outputs=target_language)
+    
+    submit_button = gr.Button("Procesar")
+    submit_button.click(process_file, inputs=[file, action, target_language], outputs=output_text)
+
+    def generate_file():
+        summary_text = output_text.value
+        filename = 'translation.txt' if action.value == 'Traducir' else 'summary.txt'
+        file_path = download_text(summary_text, filename)
+        return file_path
+
+    download_button = gr.Button("Descargar Resultado")
+    download_button.click(
+        fn=generate_file,
+        inputs=[],
+        outputs=gr.File()
+    )
+
+# Ejecutar la aplicación Gradio
+demo.launch(share=True)