Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import re
 import unicodedata
 import io
@@ -10,6 +9,7 @@ from transformers import (
     AutoTokenizer,
     AutoModelForTokenClassification,
     AutoModelForCausalLM,
+    BitsAndBytesConfig,
 )
 from peft import PeftModel
 
@@ -36,37 +36,44 @@ ID2LABEL    = {0: "O", 1: "B-TIMEX", 2: "I-TIMEX"}
 BASE_ID     = "google/gemma-2b-it"
 ADAPTER_ID  = "Rhulli/gemma-2b-it-TIMEX3"
 
-# --- Leer el token del entorno (añadido como Repository Secret) ---
-HF_TOKEN = os.getenv("HF_TOKEN")
+# --- Configuración de cuantización ---
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
 
 def load_models():
-    # Carga del modelo NER
-    ner_tok = AutoTokenizer.from_pretrained(NER_ID, token=HF_TOKEN)
-    ner_mod = AutoModelForTokenClassification.from_pretrained(NER_ID, token=HF_TOKEN)
+    # Modelo NER
+    ner_tok = AutoTokenizer.from_pretrained(NER_ID)
+    ner_mod = AutoModelForTokenClassification.from_pretrained(NER_ID)
     ner_mod.eval()
     if torch.cuda.is_available():
         ner_mod.to("cuda")
 
-    # Carga del modelo base de normalización (sin cuantización)
-    base_mod = AutoModelForCausalLM.from_pretrained(
-        BASE_ID,
-        device_map="auto",
-        token=HF_TOKEN
-    )
+    # Modelo de normalización (solo 4bit si hay GPU)
+    if torch.cuda.is_available():
+        base_mod = AutoModelForCausalLM.from_pretrained(
+            BASE_ID,
+            quantization_config=quant_config,
+            device_map="auto"
+        )
+    else:
+        base_mod = AutoModelForCausalLM.from_pretrained(
+            BASE_ID,
+            device_map="auto"
+        )
 
-    # Carga del tokenizer y adaptador LoRA
-    norm_tok = AutoTokenizer.from_pretrained(ADAPTER_ID, use_fast=True, token=HF_TOKEN)
+    norm_tok = AutoTokenizer.from_pretrained(ADAPTER_ID, use_fast=True)
     norm_mod = PeftModel.from_pretrained(
         base_mod,
         ADAPTER_ID,
-        device_map="auto",
-        token=HF_TOKEN
+        device_map="auto"
     )
     norm_mod.eval()
 
     return ner_tok, ner_mod, norm_tok, norm_mod
 
-# Carga inicial de los modelos
 ner_tok, ner_mod, norm_tok, norm_mod = load_models()
 eos_id = norm_tok.convert_tokens_to_ids("<end_of_turn>")
 
@@ -89,7 +96,7 @@ def read_file(file_obj) -> str:
         except:
             return data.decode('latin-1', errors='ignore')
 
-# --- Procesamiento de texto ---
+# --- Procesamiento ---
 def extract_timex(text: str):
     text_norm = _normalise_spaces(_normalise_apostrophes(text))
     inputs = ner_tok(text_norm, return_tensors="pt", truncation=True)
@@ -176,20 +183,69 @@ with gr.Blocks() as demo:
 
         Esta aplicación permite extraer expresiones temporales de textos o archivos (.txt, .pdf)
         y normalizarlas a formato TIMEX3.
+
+        **Cómo usar:**
+        - Sube uno o varios archivos en la columna izquierda.
+        - Ajusta la *Fecha de Anclaje (DCT)* justo debajo de los archivos.
+        - Escribe o pega tu texto en la columna derecha.
+        - Pulsa **Procesar** para ver los resultados en la tabla debajo.
+
+        **Columnas de salida:**
+        - *Expresión*: la frase temporal extraída.
+        - *Normalización*: la etiqueta TIMEX3 generada.
         """
     )
 
     with gr.Row():
         with gr.Column(scale=1):
-            files     = gr.File(file_types=['.txt', '.pdf'], file_count='multiple', label='Archivos (.txt, .pdf)')
-            dct_input = gr.Textbox(value="2025-06-11", label="Fecha de Anclaje (YYYY-MM-DD)")
+            files     = gr.File(
+                file_types=['.txt', '.pdf'],
+                file_count='multiple',
+                label='Archivos (.txt, .pdf)'
+            )
+            dct_input = gr.Textbox(
+                value="2025-06-11",
+                label="Fecha de Anclaje (YYYY-MM-DD)"
+            )
             run_btn   = gr.Button("Procesar")
+            download_btn = gr.Button("Descargar CSV")
         with gr.Column(scale=2):
-            raw_text  = gr.Textbox(lines=15, placeholder='Pega o escribe aquí tu texto...', label='Texto libre')
+            raw_text  = gr.Textbox(
+                lines=15,
+                placeholder='Pega o escribe aquí tu texto... (opcional si subes archivos)',
+                label='Texto libre'
+            )
+
+    output_table = gr.Dataframe(
+        headers=['Expresión', 'Normalización'],
+        label="Resultados",
+        interactive=False,
+        datatype=["str", "str"],
+        type="pandas"
+    )
+    output_logs  = gr.Textbox(
+        label="Logs",
+        lines=5,
+        interactive=False
+    )
+    csv_file_output = gr.File(label="Descargar resultados en CSV", visible=False)
 
-    output_table = gr.Dataframe(headers=['Expresión', 'Normalización'], label="Resultados")
-    output_logs  = gr.Textbox(label="Logs", lines=5, interactive=False)
+    run_btn.click(
+        fn=run_pipeline,
+        inputs=[files, raw_text, dct_input],
+        outputs=[output_table, output_logs]
+    )
 
-    run_btn.click(fn=run_pipeline, inputs=[files, raw_text, dct_input], outputs=[output_table, output_logs])
+    def export_csv(df):
+        csv_io = io.StringIO()
+        df.to_csv(csv_io, index=False)
+        csv_io.seek(0)
+        return gr.File.update(value=csv_io, visible=True)
+
+    download_btn.click(
+        fn=export_csv,
+        inputs=[output_table],
+        outputs=csv_file_output
+    )
 
     demo.launch()