update preprocess structure and model

tasks/text.py

@@ -4,10 +4,17 @@ from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
 from skops.io import load
+# Textpreprocessor defined in this scope
+import nltk
+
+# Download required NLTK resources
+nltk.download('punkt_tab')
+nltk.download('wordnet')
 
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
+from .utils.text_preprocessor import preprocess
 
 router = APIRouter()
 
@@ -20,8 +27,6 @@ models_description = {
 }
 
 
-
-
 # Some code borrowed from Nonnormalizable
 
 def baseline_model(dataset_length: int):
@@ -31,46 +36,15 @@ def baseline_model(dataset_length: int):
     return predictions
 
 def tree_classifier(test_dataset: dict, model: str):
-    # Textpreprocessor defined in this scope
-    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 sklearn transformer for preprocessing text
-    class TextPreprocessor(BaseEstimator, TransformerMixin):
-        def __init__(self):
-            self.lemmatizer = WordNetLemmatizer()
-
-        def fit(self, X, y=None):
-            return self
-
-        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
 
     texts = test_dataset["quote"]
 
+    texts = preprocess(texts)
+
     model_path = f"tasks/text_models/{model}.skops"
 
     model = load(model_path, 
                  trusted=[
-                     '__main__.TextPreprocessor',
-                     'nltk.stem.wordnet.WordNetLemmatizer',
                      'xgboost.core.Booster',
                      'xgboost.sklearn.XGBClassifier'])
 

tasks/text_models/xgb_pipeline.skops

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

tasks/utils/text_preprocessor.py

@@ -0,0 +1,18 @@
+import nltk
+from nltk.stem import WordNetLemmatizer
+import contractions
+
+
+
+def preprocess(X):
+    lemmatizer  = WordNetLemmatizer()
+    preprocessed_texts = []
+    for doc in X:
+        # Expand contractions
+        expanded = contractions.fix(doc)
+        # Lowercase
+        lowered = expanded.lower()
+        # Tokenize and lemmatize
+        lemmatized = " ".join([lemmatizer.lemmatize(word) for word in nltk.word_tokenize(lowered)])
+        preprocessed_texts.append(lemmatized)
+    return preprocessed_texts