updated training script

finance/train_finbert.py

@@ -11,85 +11,79 @@ from transformers import (
     AutoTokenizer,
     AutoModelForSequenceClassification,
     Trainer,
-    TrainingArguments,
+    TrainingArguments
 )
 from datasets import Dataset
 import numpy as np
 from sklearn.metrics import accuracy_score, f1_score
 import os
 
-
 def compute_metrics(p):
     """Computes and returns evaluation metrics."""
     preds = np.argmax(p.predictions, axis=1)
-    f1 = f1_score(p.label_ids, preds, average="weighted", zero_division=0)
+    f1 = f1_score(p.label_ids, preds, average='weighted', zero_division=0)
     acc = accuracy_score(p.label_ids, preds)
     return {"accuracy": acc, "f1": f1}
 
-
 def main():
     """Main function to load data, fine-tune the FinBERT model, and save it."""
-    # 1. Load and Prepare the Dataset
-    print("Loading and preparing the dataset...")
+    # --- 1. Define Paths for Colab Environment ---
+    # This path points directly to the file inside your Google Drive
+    data_path = "/content/drive/MyDrive/Colab_Data/emotion_dataset.csv"
+    # This is where the final trained model will be saved in your Google Drive
+    output_dir = "/content/drive/MyDrive/Colab_Data/finbert_emotion_model"
+    
+    # --- 2. Load and Prepare the Dataset ---
+    print(f"Loading dataset from: {data_path}")
     try:
-        # Build the correct path to the data file
-        script_dir = os.path.dirname(os.path.abspath(__file__))
-        project_root = os.path.dirname(script_dir)
-        data_path = os.path.join(project_root, "Data", "emotion_dataset.csv")
         df = pd.read_csv(data_path)
     except FileNotFoundError:
-        print(f"Error: Could not find emotion_dataset.csv at '{data_path}'")
+        print("ERROR: 'emotion_dataset.csv' not found in '/content/drive/MyDrive/Colab_Data/'.")
+        print("Please make sure you have uploaded the file to the correct Google Drive folder.")
         return
 
-    df["text"] = df["Clean_Text"].fillna(df["Text"])
-    df.dropna(subset=["text", "Emotion"], inplace=True)
-
+    print("Dataset loaded successfully. Preprocessing data...")
+    df['text'] = df['Clean_Text'].fillna(df['Text'])
+    df.dropna(subset=['text', 'Emotion'], inplace=True)
+    
     label_encoder = LabelEncoder()
-    df["labels"] = label_encoder.fit_transform(df["Emotion"])
+    df['labels'] = label_encoder.fit_transform(df['Emotion'])
     num_labels = len(label_encoder.classes_)
-
     id2label = {i: label for i, label in enumerate(label_encoder.classes_)}
     label2id = {label: i for i, label in enumerate(label_encoder.classes_)}
+    
     print(f"Found {num_labels} unique emotions.")
-
-    train_df, val_df = train_test_split(
-        df, test_size=0.2, random_state=42, stratify=df["labels"]
-    )
+    train_df, val_df = train_test_split(df, test_size=0.2, random_state=42, stratify=df['labels'])
     train_dataset = Dataset.from_pandas(train_df)
     val_dataset = Dataset.from_pandas(val_df)
 
-    # 2. Load Tokenizer and Preprocess Data 
-    print("Loading tokenizer and preprocessing data...")
+    # --- 3. Load Tokenizer and Model ---
+    print("Loading FinBERT tokenizer and model...")
     model_name = "ProsusAI/finbert"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-    def tokenize_function(examples):
-        return tokenizer(
-            examples["text"], padding="max_length", truncation=True, max_length=128
-        )
-
-    tokenized_train_dataset = train_dataset.map(tokenize_function, batched=True)
-    tokenized_val_dataset = val_dataset.map(tokenize_function, batched=True)
-
-    # 3. Load and Configure the Model 
-    print("Loading pre-trained model...")
+    
     model = AutoModelForSequenceClassification.from_pretrained(
         model_name,
         num_labels=num_labels,
         id2label=id2label,
         label2id=label2id,
-        ignore_mismatched_sizes=True, # The pre-trained model has a different number of labels
+        ignore_mismatched_sizes=True # ESSENTIAL for transfer learning
     )
+    
+    def tokenize_function(examples):
+        return tokenizer(examples['text'], padding="max_length", truncation=True, max_length=128)
 
-    # 4. Fine-Tune the Model
+    tokenized_train_dataset = train_dataset.map(tokenize_function, batched=True)
+    tokenized_val_dataset = val_dataset.map(tokenize_function, batched=True)
+
+    # --- 4. Fine-Tune the Model ---
     print("Starting model fine-tuning...")
-    output_dir = os.path.join(script_dir, "finbert_emotion_model")
     print(f"Model will be saved to: {output_dir}")
 
     training_args = TrainingArguments(
         output_dir=output_dir,
         num_train_epochs=1,
-        per_device_train_batch_size=16,
+        per_device_train_batch_size=16,  # Safe batch size for Colab GPU
         per_device_eval_batch_size=16,
         logging_steps=100,
         evaluation_strategy="epoch",
@@ -107,23 +101,14 @@ def main():
 
     trainer.train()
 
-    # 5. Save and Verify the Model
-    print(f"Training complete. Saving model and tokenizer to '{output_dir}'...")
+    # --- 5. Save the Final Model ---
+    print(f"Training complete. Saving final model to '{output_dir}'...")
     trainer.save_model(output_dir)
     tokenizer.save_pretrained(output_dir)
-    print("Model and tokenizer saved successfully!")
-    
-    # Verification step
-    try:
-        saved_files = os.listdir(output_dir)
-        print("\n--- VERIFICATION ---")
-        print(f"Successfully found the following files in '{output_dir}':")
-        for file_name in saved_files:
-            print(f"- {file_name}")
-        print("--------------------")
-    except Exception as e:
-        print(f"Could not verify saved files. Error: {e}")
-
+    print("-----" * 10)
+    print("SUCCESS: Model saved to your Google Drive in the 'Colab_Data' folder.")
+    print("You can now download the 'finbert_emotion_model' folder and place it in your project's 'finance/' directory.")
+    print("-----" * 10)
 
 if __name__ == "__main__":
     main()