Spaces:
Running
Running
File size: 3,510 Bytes
e062e72 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from datasets import load_dataset
import joblib
import os
import numpy as np
# Define paths for the Random Forest model, TF-IDF vectorizer, and label encoder
rf_model_path = 'random_forest_model.pkl'
vectorizer_path = "tfidf_vectorizer.pkl"
label_encoder_path = "label_encoder.pkl"
multi_rf_model_path= "random_forest_multi_model.pkl"
# Check if models and encoder exist
if os.path.exists(rf_model_path) and os.path.exists(vectorizer_path) and os.path.exists(label_encoder_path) and os.path.exists(multi_rf_model_path):
# Load the models if they already exist
rf_single = joblib.load(rf_model_path)
vectorizer = joblib.load(vectorizer_path)
le = joblib.load(label_encoder_path)
rf_multi = joblib.load(multi_rf_model_path)
print("Random Forest model, vectorizer, and label encoder loaded from disk.")
else:
# Load the dataset
ds = load_dataset('ahmedheakl/resume-atlas', cache_dir="C:/Users/dell/.cache/huggingface/datasets")
# Create a DataFrame from the 'train' split
df_train = pd.DataFrame(ds['train'])
# Initialize the Label Encoder and encode the 'Category' labels
le = LabelEncoder()
df_train['Category_encoded'] = le.fit_transform(df_train['Category'])
# Split the dataset into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
df_train['Text'], df_train['Category_encoded'], test_size=0.2, random_state=42)
# Initialize TF-IDF Vectorizer and transform the text data
vectorizer = TfidfVectorizer(max_features=1000)
X_train_tfidf = vectorizer.fit_transform(X_train)
X_test_tfidf = vectorizer.transform(X_test)
# Initialize and train the Random Forest models
rf_single = RandomForestClassifier(n_estimators=100, random_state=42)
rf_single.fit(X_train_tfidf, y_train)
rf_multi = RandomForestClassifier(n_estimators=100, random_state=42)
rf_multi.fit(X_train_tfidf, y_train)
# Save the Random Forest models, TF-IDF vectorizer, and label encoder
joblib.dump(rf_single, rf_model_path)
joblib.dump(rf_multi, multi_rf_model_path)
joblib.dump(vectorizer, vectorizer_path)
joblib.dump(le, label_encoder_path)
print("Random Forest model, vectorizer, and label encoder trained and saved to disk.")
# Single-label classification function for Random Forest model
def classify_text_rf(text):
try:
text_tfidf = vectorizer.transform([text])
predicted_class_index = rf_single.predict(text_tfidf)[0]
predicted_category = le.inverse_transform([predicted_class_index])[0]
return predicted_category
except Exception as e:
print(f"Error in classify_text_rf: {e}")
return None
# Multi-label classification function with top N predictions
def classify_text_rf_multi(text, top_n=3):
try:
text_tfidf = vectorizer.transform([text])
probabilities = rf_multi.predict_proba(text_tfidf)[0]
top_n_indices = np.argsort(probabilities)[::-1][:min(top_n, len(probabilities))]
top_n_categories = le.inverse_transform(top_n_indices)
return top_n_categories
except Exception as e:
print(f"Error in classify_text_rf_multi: {e}")
return None
|