Update app.py

app.py

@@ -1,11 +1,11 @@
-import gradio as gr
+    import gradio as gr
 import json
-from span_marker import SpanMarkerModel, SpanMarkerTrainer, SpanMarkerTrainingArguments
+from span_marker import SpanMarkerModel, SpanMarkerTrainer
+from span_marker import SpanMarkerTrainingArguments
 from datasets import Dataset
 from sklearn.model_selection import train_test_split
 
 def entrenar(jsonl_file):
-    # Cargar JSONL
     raw = [json.loads(l) for l in jsonl_file.splitlines()]
 
     dataset = []
@@ -25,9 +25,9 @@ def entrenar(jsonl_file):
 
     # Extraer etiquetas
     labels = sorted(list({e["label"] for d in dataset for e in d["entities"]}))
-    labels.insert(0, "O")  # obligatorio
+    labels.insert(0, "O")
 
-    # Train/test
+    # Datasets Hugging Face
     train, test = train_test_split(dataset, test_size=0.2, random_state=42)
     train_ds = Dataset.from_list(train)
     test_ds = Dataset.from_list(test)
@@ -38,19 +38,18 @@ def entrenar(jsonl_file):
         labels=labels
     )
 
-    # Argumentos
+    # Args
     args = SpanMarkerTrainingArguments(
         output_dir="modelo_final",
+        num_train_epochs=3,
         learning_rate=5e-5,
         per_device_train_batch_size=2,
         per_device_eval_batch_size=2,
-        num_train_epochs=3,
-        logging_steps=10,
         save_strategy="epoch",
-        evaluation_strategy="epoch"
+        evaluation_strategy="epoch",
+        logging_steps=10
     )
 
-    # Entrenador
     trainer = SpanMarkerTrainer(
         model=model,
         args=args,
@@ -60,13 +59,14 @@ def entrenar(jsonl_file):
 
     trainer.train()
 
-    return "Entrenamiento completado. El modelo está en /modelo_final"
+    return "¡Entrenamiento completado! Modelo guardado en /modelo_final"
 
 ui = gr.Interface(
     fn=entrenar,
-    inputs=gr.File(label="Sube tu archivo JSONL exportado de Label Studio"),
+    inputs=gr.File(label="Sube tu JSONL exportado de Label Studio"),
     outputs="text",
     title="Entrenamiento NER Médico con SpanMarker"
 )
 
 ui.launch()
+