import os.path import pickle import pandas as pd from transformers import AutoTokenizer from sklearn.model_selection import train_test_split from sklearn.feature_extraction.text import CountVectorizer from xgboost import XGBClassifier from sklearn.preprocessing import LabelEncoder class MentalHealthClassifier: def __init__(self, data_path, model_path): self.data = pd.read_csv(data_path, skip_blank_lines=True) self.data['category'] = ['anxiety' if x == 'axienty' else x for x in self.data['category']] self.data.dropna(subset=['text'], inplace=True) self.data.dropna(subset=['clean_text'], inplace=True) self.data_selected = self.data[['clean_text', 'category']] self.df = pd.DataFrame(self.data_selected) self.label_encoder = LabelEncoder() self.df['category_encoded'] = self.label_encoder.fit_transform(self.df['category']) self.tokenizer = None self.vectorizer = CountVectorizer() self.model_path = model_path self.model = self.load_model() if os.path.exists(model_path) else XGBClassifier() def preprocess_data(self): tokenized_texts = [self.tokenizer.tokenize(text, padding=True, truncation=True) for text in self.df['clean_text']] X = self.vectorizer.fit_transform([' '.join(tokens) for tokens in tokenized_texts]).toarray() return X, self.df['category_encoded'] def train_model(self, X, y): X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) self.model.fit(X_train, y_train) y_pred = self.model.predict(X_test) return y_test, y_pred def predict_category(self, input_text): if self.tokenizer is None: raise ValueError("Tokenizer not initialized. Call 'initialize_tokenizer' first.") tokenized_input = self.tokenizer.tokenize(input_text, padding=True, truncation=True) input_feature_vector = self.vectorizer.transform([' '.join(tokenized_input)]).toarray() predicted_category_encoded = self.model.predict(input_feature_vector) predicted_category = self.label_encoder.inverse_transform(predicted_category_encoded) return predicted_category[0] def initialize_tokenizer(self, model_name): self.model_name = model_name self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) def save_model(self): print("saving model...to pickle...") with open(self.model_path, 'wb') as f: pickle.dump(self.model, f) def load_model(self): print("loading model...from pickle...") with open(self.model_path, 'rb') as f: return pickle.load(f) if __name__ == "__main__": tokenizer_model_name = "nlptown/bert-base-multilingual-uncased-sentiment" data_path = 'data/data.csv' model_path = 'mental_health_model.pkl' mental_classifier = MentalHealthClassifier(data_path, model_path) if not os.path.exists(model_path): mental_classifier.initialize_tokenizer(tokenizer_model_name) X, y = mental_classifier.preprocess_data() y_test, y_pred = mental_classifier.train_model(X, y) mental_classifier.save_model() else: mental_classifier.load_model(model_path) mental_classifier.initialize_tokenizer(tokenizer_model_name) # Ensure tokenizer is initialized if loading model from pickle mental_classifier.preprocess_data() # input_text = "I feel anxiety whenever i am doing nothing." # predicted_category = mental_classifier.predict_category(input_text) # print("Predicted mental health condition:", predicted_category)