Update tasks/text.py

tasks/text.py

@@ -12,6 +12,22 @@ import joblib
 REPO_ID = "kantundpeterpan/frugal-ai-toy"
 FILENAME = "tfidf_rf.pkl"
 
+import nltk
+from nltk.tokenize import WordPunctTokenizer
+from nltk.stem import WordNetLemmatizer
+from nltk.corpus import stopwords
+import string
+
+stop = set(stopwords.words('english') + list(string.punctuation))
+
+def tokenize_quote(r):
+    tokens = nltk.word_tokenize(r.lower())
+    cleaned = [word for word in tokens if word not in stop]
+    return cleaned
+
+def lemmatize_tokens(tokens: list):
+    return [lemmatizer.lemmatize(t) for t in tokens]
+
 def lemmatize_X(X):
     return X.quote.apply(tokenize_quote).apply(lemmatize_tokens).apply(lambda x: " ".join(x))