adding training for new model

dillon · dillon · commit ef85461cbe1a · 2025-07-31T12:01:35.000+02:00
diff --git a/data/processed/bow_best_params.json b/data/processed/bow_best_params.json
@@ -0,0 +1 @@
+{}
diff --git a/data/processed/tfidf_clf_params.json b/data/processed/tfidf_clf_params.json
@@ -0,0 +1,7 @@
+{
+  "C": 1,
+  "class_weight": "balanced",
+  "max_iter": 500,
+  "penalty": "l2",
+  "solver": "lbfgs"
+}
diff --git a/main.py b/main.py
@@ -1,19 +1,22 @@
 from src.data.processor import unzip_data_extract_contents
-from src.model.train_model import train_bow_logreg, train_tfidf_logreg
-from src.model.predict_model import predict_sentiment_tfidf, predict_sentiment_bow
+from src.model.train_model import train_bow_logreg, train_tfidf_logreg, finetune_bert
+from src.model.predict_model import predict_sentiment_tfidf, predict_sentiment_bow, predict_sentiment_bert
 
 def main():
     test_archive, train_archive, unsup_archive, imdb_vocab, imdb_expected_rating = unzip_data_extract_contents()
     train_tfidf_logreg(test_archive, train_archive, imdb_vocab, imdb_expected_rating)
     train_bow_logreg(test_archive, train_archive, imdb_vocab)
+    finetune_bert(train_archive, test_archive)
 
     sample_text = "this movie was mid"
     print("TFIDF Prediction:", predict_sentiment_tfidf(sample_text))
     print("BoW Prediction:", predict_sentiment_bow(sample_text, imdb_vocab))
+    print("BERT Prediction:", predict_sentiment_bert(sample_text))
 
     sample_text = "it wasnt bad"
     print("TFIDF Prediction:", predict_sentiment_tfidf(sample_text))
     print("BoW Prediction:", predict_sentiment_bow(sample_text, imdb_vocab))
+    print("BERT Prediction:", predict_sentiment_bert(sample_text))
 
 
 main()
diff --git a/models/bow_sentiment_model.joblib b/models/bow_sentiment_model.joblib
diff --git a/models/bow_sentiment_model_with_params.joblib b/models/bow_sentiment_model_with_params.joblib
diff --git a/models/tfidf_sentiment_model_with_params.joblib b/models/tfidf_sentiment_model_with_params.joblib
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,3 @@
-transformers==4.53.3
 jupyter==1.1.1
 numpy==2.3.1
 matplotlib==3.10.3
@@ -7,4 +6,8 @@ joblib==1.5.1
 flake8==7.0.0
 pre-commit==4.2.0
 pytest==8.4.1
-coverage==7.4.4
+coverage==7.4.4
+torch==2.7.1
+transformers==4.54.1
+datasets==4.0.0
+transformers[torch]==4.54.1
diff --git a/src/common/utils.py b/src/common/utils.py
@@ -4,84 +4,61 @@
 import joblib
 from scipy.sparse import save_npz, load_npz
 
+
+def get_project_root():
+    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def get_processed_path(file_name, is_json=False):
+    ext = '.json' if is_json else '.gz'
+    if not file_name.endswith(ext):
+        file_name += ext
+    return os.path.join(get_project_root(), "data", "processed", file_name)
+
+
 def export_data_to_json(data, file_name, is_json=False):
-    project_root = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
+    path = get_processed_path(file_name, is_json)
     if is_json:
-        if not file_name.endswith('.json'):
-            file_name += '.json'
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
     else:
-        if not file_name.endswith('.gz'):
-            file_name += '.gz'
-    path = os.path.join(project_root, "data", "processed", file_name)
+        with gzip.open(path, 'wt', encoding='utf-8') as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
 
-    with gzip.open(path, 'wt', encoding='utf-8') as f:
-        json.dump(data, f, ensure_ascii=False, indent=2)
-        
 
 def import_processed_json(file_name, is_json=False):
-    project_root = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
+    path = get_processed_path(file_name, is_json)
+    if not os.path.exists(path):
+        return None
     if is_json:
-        if not file_name.endswith('.json'):
-            file_name += '.json'
-        path = os.path.join(project_root, "data", "processed", file_name)
-        if not os.path.exists(path):
-            return None
         with open(path, 'r', encoding='utf-8') as f:
-            data = json.load(f)
-        return data
+            return json.load(f)
     else:
-        if not file_name.endswith('.gz'):
-            file_name += '.gz'
-        path = os.path.join(project_root, "data", "processed", file_name)
-        if not os.path.exists(path):
-            return None
         with gzip.open(path, 'rt', encoding='utf-8') as f:
-            data = json.load(f)
-        return data
+            return json.load(f)
 
-def export_models(data, file_name):
-    project_root = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
-    path = os.path.join(project_root, "models", file_name)
 
-    import joblib
-    joblib.dump(data, path)
+def export_models(model, file_name):
+    path = os.path.join(get_project_root(), "models", file_name)
+    joblib.dump(model, path)
 
 
 def import_models(file_name):
-    project_root = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
-    path = os.path.join(project_root, "models", file_name)
-
+    path = os.path.join(get_project_root(), "models", file_name)
     if not os.path.exists(path):
         return None
-
-    data = joblib.load(path)
-    return data
+    return joblib.load(path)
 
 
 def export_processed_data(matrix, filename):
-    project_root = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
-    processed_dir = os.path.join(project_root, "data", "processed")
+    processed_dir = os.path.join(get_project_root(), "data", "processed")
     os.makedirs(processed_dir, exist_ok=True)
-
     path = os.path.join(processed_dir, filename)
     save_npz(path, matrix)
 
 
 def import_processed_data(filename):
-    project_root = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
-    processed_dir = os.path.join(project_root, "data", "processed")
+    processed_dir = os.path.join(get_project_root(), "data", "processed")
     path = os.path.join(processed_dir, filename)
     if not os.path.exists(path):
         return None
diff --git a/src/data/processor.py b/src/data/processor.py
@@ -7,6 +7,7 @@
 import html
 import numpy as np
 from scipy.sparse import lil_matrix
+from datasets import Dataset
 
 def unzip_data_extract_contents():
     project_root = os.path.dirname(
@@ -142,4 +143,13 @@ def parse_bow_line(line):
         if ':' in part:
             idx, val = part.split(':')
             bow[int(idx)] = int(val)
-    return bow
+    return bow
+
+def prepare_bert_dataset(archive):
+    texts = [r['contents'] for r in archive.reviews]
+    labels = [0 if r['type'] == 'neg' else 1 for r in archive.reviews]
+    return Dataset.from_dict({"text": texts, "label": labels})
+
+
+def tokenize_function(examples, tokenizer):
+    return tokenizer(examples["text"], truncation=True, padding="max_length")
diff --git a/src/model/predict_model.py b/src/model/predict_model.py
@@ -1,6 +1,8 @@
 from src.common.utils import import_models
 from src.data.processor import clean_review_text, bow_dicts_to_matrix
 from src.common.utils import import_models
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
 
 def predict_sentiment_tfidf(text):
@@ -35,4 +37,15 @@ def text_to_bow_dict(text, vocab_list):
         idx = vocab_index.get(word)
         if idx is not None:
             bow[idx] = bow.get(idx, 0) + 1
-    return bow
+    return bow
+
+
+def predict_sentiment_bert(text, model_dir="./bert_finetuned"):
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        sentiment = torch.argmax(probs, dim=1).item()
+    return sentiment, probs.squeeze().tolist()
diff --git a/src/model/train_model.py b/src/model/train_model.py
@@ -1,11 +1,15 @@
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
 from sklearn.metrics import classification_report
-from src.common.utils import export_models, import_models, import_processed_json, export_data_to_json
+from src.common.utils import export_models, import_models, import_processed_json, export_data_to_json, get_project_root
 from src.data.processor import bow_dicts_to_matrix, parse_bow_line
 from scipy.sparse import csr_matrix
 from sklearn.model_selection import GridSearchCV
 from sklearn.pipeline import Pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+from datasets import Dataset
+import torch
+import os
 
 # Function to train a TF-IDF model using Logistic Regression
 def train_tfidf_logreg(test_data, train_data, vocab_list, imdb_expected_rating=None):
@@ -36,9 +40,8 @@ def train_tfidf_logreg(test_data, train_data, vocab_list, imdb_expected_rating=N
 
     clf_with_best_params = import_models('tfidf_sentiment_model_with_params.joblib')
     if clf_with_best_params is None:
-        # Pass X_train (list of strings), not X_train_vec
         best_params = tune_tfidf_logreg(X_train, y_train, vocab_list)
-        clf_with_best_params = LogisticRegression(**best_params, max_iter=1000)
+        clf_with_best_params = LogisticRegression(**best_params)
         clf_with_best_params.fit(X_train_vec, y_train)
         export_models(clf_with_best_params, 'tfidf_sentiment_model_with_params.joblib')
 
@@ -78,7 +81,7 @@ def train_bow_logreg(test_archive, train_archive, vocab_list):
     clf_with_best_params = import_models('bow_sentiment_model_with_params.joblib')
     if clf_with_best_params is None:
         best_params = tune_bow_logreg(X_train, y_train)
-        clf_with_best_params = LogisticRegression(**best_params, max_iter=1000)
+        clf_with_best_params = LogisticRegression(**best_params)
         clf_with_best_params.fit(X_train, y_train)
         export_models(clf_with_best_params, 'bow_sentiment_model_with_params.joblib')
     y_pred = clf.predict(X_test)
@@ -96,18 +99,15 @@ def tune_tfidf_logreg(X_train, y_train, vocab_list):
         ('clf', LogisticRegression(max_iter=1000))
     ])
     param_grid = {
-        'tfidf__ngram_range': [(1,1), (1,2), (1,3), (2,2), (2,3)],
-        'tfidf__max_features': [1000, 5000, 10000, 20000, 50000],
-        'tfidf__min_df': [1, 2, 5, 10],
-        'tfidf__max_df': [0.7, 0.8, 0.9, 1.0],
+        'tfidf__ngram_range': [(1,1), (1,2), (1,3)],
+        'tfidf__max_features': [5000, 10000, 20000],
+        'tfidf__min_df': [1, 2, 5],
         'tfidf__stop_words': [None, 'english'],
-        'tfidf__sublinear_tf': [True, False],
-        'tfidf__norm': ['l1', 'l2'],
-        'clf__C': [0.01, 0.1, 1, 10, 100],
-        'clf__solver': ['lbfgs', 'liblinear', 'saga'],
-        'clf__penalty': ['l2', 'none'],
+        'clf__C': [0.1, 1, 10],
+        'clf__solver': ['lbfgs', 'liblinear'],
+        'clf__penalty': ['l2'],
         'clf__class_weight': [None, 'balanced'],
-        'clf__max_iter': [500, 1000, 2000]
+        'clf__max_iter': [500, 1000]
     }
     grid = GridSearchCV(pipeline, param_grid, cv=3, n_jobs=-1, verbose=1)
     grid.fit(X_train, y_train)
@@ -124,14 +124,14 @@ def tune_bow_logreg(X_train, y_train):
         return best_params
 
     param_grid = {
-        'C': [0.001, 0.01, 0.1, 1, 10, 100],
-        'max_iter': [100, 500, 1000, 2000],
-        'solver': ['lbfgs', 'liblinear', 'saga', 'newton-cg'],
-        'penalty': ['l1', 'l2', 'none'],
+        'C': [0.1, 1, 10],
+        'max_iter': [500, 1000],
+        'solver': ['lbfgs', 'liblinear'],
+        'penalty': ['l2'],
         'class_weight': [None, 'balanced'],
-        'fit_intercept': [True, False],
-        'warm_start': [True, False],
-        'tol': [1e-4, 1e-3, 1e-2]
+        'fit_intercept': [True],
+        'warm_start': [False],
+        'tol': [1e-4, 1e-3]
     }
     clf = LogisticRegression()
     grid = GridSearchCV(clf, param_grid, cv=3, n_jobs=-1, verbose=1)
@@ -140,4 +140,59 @@ def tune_bow_logreg(X_train, y_train):
     bow_best_params = {k.replace('clf__', ''): v for k, v in grid.best_params_.items() if k.startswith('clf__')}
 
     export_data_to_json(bow_best_params, 'bow_best_params', is_json=True)
-    return bow_best_params
+    return bow_best_params
+
+
+def prepare_bert_dataset(archive):
+    texts = [r['contents'] for r in archive.reviews]
+    labels = [0 if r['type'] == 'neg' else 1 for r in archive.reviews]
+    return Dataset.from_dict({"text": texts, "label": labels})
+
+
+def tokenize_function(examples, tokenizer):
+    return tokenizer(examples["text"], truncation=True, padding="max_length")
+
+
+def import_finetuned_bert():
+    model_dir = os.path.join(get_project_root(), "models", "bert_finetuned")
+    if not os.path.exists(model_dir):
+        return None, None
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+    return model, tokenizer
+
+def finetune_bert(train_archive, test_archive, model_name="bert-base-uncased"):
+    model, tokenizer = import_finetuned_bert()
+    if model is None or tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        train_dataset = prepare_bert_dataset(train_archive)
+        test_dataset = prepare_bert_dataset(test_archive)
+
+        train_dataset = train_dataset.map(lambda x: tokenize_function(x, tokenizer), batched=True)
+        test_dataset = test_dataset.map(lambda x: tokenize_function(x, tokenizer), batched=True)
+        train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+        test_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+
+        model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
+
+        bert_model_dir = os.path.join(get_project_root(), "models", "bert_finetuned")
+
+        training_args = TrainingArguments(
+            output_dir=bert_model_dir,
+            num_train_epochs=2,
+            per_device_train_batch_size=8,
+            logging_steps=50,
+            save_steps=100,
+            report_to=[],
+        )
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=test_dataset,
+        )
+
+        trainer.train()
+        model.save_pretrained(bert_model_dir)
+        tokenizer.save_pretrained(bert_model_dir)
+    return model, tokenizer
diff --git a/train.py b/train.py
