update code to include tree classifier

tasks/text.py

@@ -3,8 +3,8 @@ from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
-from skops.hub_utils import download
 from skops.io import load
+from .utils.text_preprocessor import TextPreprocessor
 
 
 from .utils.evaluation import TextEvaluationRequest
@@ -15,12 +15,40 @@ router = APIRouter()
 DESCRIPTION = "Random Baseline"
 ROUTE = "/text"
 
-MODEL_PATH = "tasks/text_models/xgb_pipeline.skops"
+models_description = {
+    "baseline": "random baseline",
+    "tfidf_xgb": "TF-IDF vectorizer and XGBoost classifier",
+}
+
+# Some code borrowed from Nonnormalizable
+
+def baseline_model(dataset_length: int):
+    # Make random predictions (placeholder for actual model inference)
+    predictions = [random.randint(0, 7) for _ in range(dataset_length)]
+
+    return predictions
+
+def tree_classifier(test_dataset: dict, model: str):
+    texts = test_dataset["quote"]
+
+    model_path = f"models/frugalai_{model}"
+
+    model = load(model_path, 
+                 trusted=[
+                     '__main__.TextPreprocessor',
+                     'nltk.stem.wordnet.WordNetLemmatizer',
+                     'xgboost.core.Booster',
+                     'xgboost.sklearn.XGBClassifier'])
+
+    predictions = model.predict(texts)
+
+    return predictions
 
 
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
-async def evaluate_text(request: TextEvaluationRequest):
+async def evaluate_text(request: TextEvaluationRequest,
+                        model: str = "baseline"):
     """
     Evaluate text classification for climate disinformation detection.
     
@@ -65,8 +93,10 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
 
-    model = load(MODEL_PATH)
-    predictions = model.predict(test_dataset["text"])
+    if model == "baseline":
+        predictions = baseline_model(len(true_labels))
+    elif model == "tfidf_xgb":
+        predictions = tree_classifier(test_dataset, model='tfidf_xgb')
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE

tasks/text_models/.gitattributes

@@ -1 +0,0 @@
-xgb_pipeline.skops filter=lfs diff=lfs merge=lfs -text

tasks/text_models/xgb_pipeline.skops

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

tasks/utils/text_preprocessor.py

@@ -0,0 +1,30 @@
+import nltk
+from nltk.stem import WordNetLemmatizer
+from sklearn.base import BaseEstimator, TransformerMixin
+import nltk
+import contractions
+
+# Download required NLTK resources
+nltk.download('punkt_tab')
+nltk.download('wordnet')
+
+# Custom transformer for preprocessing text
+class TextPreprocessor(BaseEstimator, TransformerMixin):
+    def __init__(self):
+        self.lemmatizer = WordNetLemmatizer()
+
+    def fit(self, X, y=None):
+        return self  # Does nothing, just returns the instance
+
+    def transform(self, X):
+        preprocessed_texts = []
+        for doc in X:
+            # Expand contractions
+            expanded = contractions.fix(doc)
+            # Lowercase
+            lowered = expanded.lower()
+
+            # Tokenize and lemmatize
+            lemmatized = " ".join([self.lemmatizer.lemmatize(word) for word in nltk.word_tokenize(lowered)])
+            preprocessed_texts.append(lemmatized)
+        return preprocessed_texts