Train LLM response quality classifier

.gitignore

@@ -0,0 +1,4 @@
+.venv/
+__pycache__/
+*.pyc
+.DS_Store

model.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a23436e774c50aa25817e699197656880faaf46ab78c2f3847f5620db38134f4
+size 44164

train.py

@@ -5,11 +5,10 @@ import joblib
 import pandas as pd
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
-from sklearn.metrics import classification_report
-from sklearn.model_selection import LeaveOneOut, cross_val_predict
+from sklearn.metrics import classification_report, confusion_matrix
+from sklearn.model_selection import StratifiedKFold, cross_val_predict
 from sklearn.pipeline import Pipeline
 
-
 DATA_PATH = Path("data/train.jsonl")
 MODEL_PATH = Path("model.joblib")
 
@@ -19,40 +18,29 @@ def load_data(path: Path) -> pd.DataFrame:
 
     with path.open("r", encoding="utf-8") as file:
         for line in file:
-            records.append(json.loads(line))
+            if line.strip():
+                records.append(json.loads(line))
 
     return pd.DataFrame(records)
 
 
-def main():
-    df = load_data(DATA_PATH)
-
-    # Combine the original prompt and response so that the classifier
-    # can use information from both parts of the evaluation example.
-    X = (
-        "PROMPT: "
-        + df["prompt"].astype(str)
-        + "\nRESPONSE: "
-        + df["response"].astype(str)
-    )
-
-    y = df["quality_label"]
-
-    pipeline = Pipeline(
+def build_pipeline() -> Pipeline:
+    return Pipeline(
         [
             (
                 "tfidf",
                 TfidfVectorizer(
                     ngram_range=(1, 2),
                     lowercase=True,
-                    min_df=1,
+                    min_df=2,
+                    max_df=0.95,
                     sublinear_tf=True,
                 ),
             ),
             (
                 "classifier",
                 LogisticRegression(
-                    max_iter=1000,
+                    max_iter=2000,
                     class_weight="balanced",
                     random_state=42,
                 ),
@@ -60,19 +48,40 @@ def main():
         ]
     )
 
-    # With only 24 examples, a conventional train/test split would
-    # provide an unstable and potentially misleading estimate.
-    loo = LeaveOneOut()
+
+def main():
+    df = load_data(DATA_PATH)
+
+    X = (
+        "PROMPT: "
+        + df["prompt"].astype(str)
+        + "\nRESPONSE: "
+        + df["response"].astype(str)
+    )
+
+    y = df["quality_label"].astype(str)
+
+    print(f"Examples: {len(df)}")
+    print(y.value_counts().sort_index())
+    print()
+
+    cv = StratifiedKFold(
+        n_splits=5,
+        shuffle=True,
+        random_state=42,
+    )
+
+    pipeline = build_pipeline()
 
     predictions = cross_val_predict(
         pipeline,
         X,
         y,
-        cv=loo,
+        cv=cv,
     )
 
-    print("Leave-One-Out Cross-Validation")
-    print("=" * 40)
+    print("Stratified 5-Fold Cross-Validation")
+    print("=" * 44)
 
     print(
         classification_report(
@@ -83,13 +92,28 @@ def main():
         )
     )
 
-    # Train the final demonstration model on all available examples.
+    labels = sorted(y.unique())
+
+    matrix = confusion_matrix(
+        y,
+        predictions,
+        labels=labels,
+    )
+
+    print("Confusion matrix")
+    print(f"Labels: {labels}")
+    print(matrix)
+    print()
+
     pipeline.fit(X, y)
 
-    joblib.dump(pipeline, MODEL_PATH)
+    joblib.dump(
+        pipeline,
+        MODEL_PATH,
+    )
 
     print(f"Saved model to: {MODEL_PATH.resolve()}")
 
 
 if __name__ == "__main__":
-    main()
+    main()