Update src/model.py

src/model.py

@@ -1,44 +1,34 @@
 # src/model.py
 from transformers import DistilBertForSequenceClassification, Trainer, TrainingArguments
 from sklearn.metrics import classification_report, confusion_matrix
-import torch
 import numpy as np
-import pandas as pd
 import logging
-from src.config import MODEL_PATH, BATCH_SIZE, EPOCHS
+from huggingface_hub import login
+from src.config import MODEL_NAME, HF_MODEL_PATH, LOCAL_MODEL_PATH, BATCH_SIZE, EPOCHS, HF_TOKEN, LOG_FILE
 
 def setup_logging():
-    logging.basicConfig(filename="logs/app.log", level=logging.INFO, 
+    logging.basicConfig(filename=LOG_FILE, level=logging.INFO, 
                         format="%(asctime)s - %(levelname)s - %(message)s")
 
-class EcommerceDataset(torch.utils.data.Dataset):
-    def __init__(self, encodings, labels):
-        self.encodings = encodings
-        self.labels = labels
+def compute_metrics(eval_pred):
+    """Compute evaluation metrics."""
+    logits, labels = eval_pred
+    predictions = np.argmax(logits, axis=-1)
+    report = classification_report(labels, predictions, output_dict=True, 
+                                  target_names=["Electronics", "Household", "Books", "Clothing & Accessories"])
+    return {"accuracy": report["accuracy"], "f1": report["macro avg"]["f1-score"]}
 
-    def __getitem__(self, idx):
-        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
-        item["labels"] = torch.tensor(self.labels[idx])
-        return item
-
-    def __len__(self):
-        return len(self.labels)
-
-def train_model(train_encodings, train_labels, val_encodings, val_labels):
-    """Fine-tune DistilBERT for classification."""
+def train_model(train_dataset, val_dataset):
+    """Fine-tune DistilBERT and push to Hugging Face Hub."""
     setup_logging()
-    model = DistilBertForSequenceClassification.from_pretrained(
-        "distilbert-base-uncased", num_labels=4
-    )
+    login(token=HF_TOKEN)  # Log in to Hugging Face Hub
+    model = DistilBertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=4)
     label_map = {"Electronics": 0, "Household": 1, "Books": 2, "Clothing & Accessories": 3}
-    train_labels = [label_map[label] for label in train_labels]
-    val_labels = [label_map[label] for label in val_labels]
-
-    train_dataset = EcommerceDataset(train_encodings, train_labels)
-    val_dataset = EcommerceDataset(val_encodings, val_labels)
+    train_dataset = train_dataset.map(lambda x: {"labels": label_map[x["category"]]})
+    val_dataset = val_dataset.map(lambda x: {"labels": label_map[x["category"]]})
 
     training_args = TrainingArguments(
-        output_dir=MODEL_PATH,
+        output_dir=LOCAL_MODEL_PATH,
         num_train_epochs=EPOCHS,
         per_device_train_batch_size=BATCH_SIZE,
         per_device_eval_batch_size=BATCH_SIZE,
@@ -46,6 +36,8 @@ def train_model(train_encodings, train_labels, val_encodings, val_labels):
         save_strategy="epoch",
         logging_dir="logs/",
         logging_steps=100,
+        push_to_hub=True,
+        hub_model_id=HF_MODEL_PATH,
     )
 
     trainer = Trainer(
@@ -53,26 +45,29 @@ def train_model(train_encodings, train_labels, val_encodings, val_labels):
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=val_dataset,
+        compute_metrics=compute_metrics,
     )
 
     logging.info("Starting model training")
     trainer.train()
-    model.save_pretrained(MODEL_PATH)
-    logging.info(f"Model saved to {MODEL_PATH}")
+    trainer.push_to_hub()  # Push model to Hugging Face Hub
+    model.save_pretrained(LOCAL_MODEL_PATH)
+    logging.info(f"Model saved locally to {LOCAL_MODEL_PATH} and pushed to {HF_MODEL_PATH}")
     return model, label_map
 
-def evaluate_model(model, test_encodings, test_labels):
+def evaluate_model(model, test_dataset):
     """Evaluate model and log metrics."""
     setup_logging()
     label_map = {"Electronics": 0, "Household": 1, "Books": 2, "Clothing & Accessories": 3}
-    test_labels = [label_map[label] for label in test_labels]
-    test_dataset = EcommerceDataset(test_encodings, test_labels)
-    trainer = Trainer(model=model)
+    test_dataset = test_dataset.map(lambda x: {"labels": label_map[x["category"]]})
+    trainer = Trainer(model=model, compute_metrics=compute_metrics)
+    results = trainer.evaluate(test_dataset)
     predictions = trainer.predict(test_dataset).predictions
     pred_labels = np.argmax(predictions, axis=1)
+    true_labels = [x["labels"] for x in test_dataset]
 
-    report = classification_report(test_labels, pred_labels, target_names=label_map.keys())
+    report = classification_report(true_labels, pred_labels, target_names=label_map.keys())
+    cm = confusion_matrix(true_labels, pred_labels)
     logging.info(f"Classification Report:\n{report}")
-    cm = confusion_matrix(test_labels, pred_labels)
     logging.info(f"Confusion Matrix:\n{cm}")
-    return report, cm
+    return report, cm, results