Update app.py (#14)

- Update app.py (ad1bd14c666771191de9e043dc62c8cd10c6c77f)

app.py

@@ -50,19 +50,51 @@ def log_event(event):
 # =========================
 # Training Pipeline
 # =========================
-def train_model(model_name, dataset_name, epochs, output_dir="trained_models"):
-    log_event(f"Training started: model={model_name}, dataset={dataset_name}, epochs={epochs}")
-
-    # Load tokenizer + dataset
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    dataset = load_dataset(dataset_name, split="train[:200]")  # smaller subset for CPU
+def train_model(model_name, dataset_name, epochs):
+    try:
+        log.info(f"Loading dataset: {dataset_name}")
+        parts = dataset_name.split(" ")
+        
+        if len(parts) == 2:
+            dataset_repo, dataset_config = parts
+            dataset = load_dataset(dataset_repo, dataset_config, split="train[:200]")  # CPU-friendly subset
+        else:
+            dataset = load_dataset(dataset_name, split="train[:200]")
+
+        log.info("Dataset loaded successfully")
+
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
+
+        def tokenize_fn(examples):
+            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
+
+        dataset = dataset.map(tokenize_fn, batched=True)
+        dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+
+        training_args = TrainingArguments(
+            output_dir="./results",
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            learning_rate=2e-5,
+            per_device_train_batch_size=4,
+            num_train_epochs=int(epochs),
+            logging_dir="./logs",
+            logging_steps=10,
+        )
 
-    def tokenize_fn(batch):
-        return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=128)
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=dataset,
+            tokenizer=tokenizer,
+        )
 
-    dataset = dataset.map(tokenize_fn, batched=True)
-    dataset = dataset.rename_column("label", "labels")
-    dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
+        trainer.train()
+        return "Training complete ✅"
+    except Exception as e:
+        log.error(f"Training failed: {e}")
+        return f"Error during training: {e}"
 
     # Load model
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)