Spaces:
Running
Running
File size: 2,811 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 |
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.preprocessing import LabelEncoder
import numpy as np
from datasets import load_dataset
import joblib
import os
# Define paths for the model, vectorizer, and label encoder
svm_model_path = "svm_resume_model.pkl"
vectorizer_path = "tfidf_vectorizer.pkl"
label_encoder_path = "label_encoder.pkl"
# Check if models exist and load them; otherwise, train and save
if os.path.exists(svm_model_path) and os.path.exists(vectorizer_path) and os.path.exists(label_encoder_path):
# Load the models if they already exist
svm_model = joblib.load(svm_model_path)
vectorizer = joblib.load(vectorizer_path)
le = joblib.load(label_encoder_path)
print("Models 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 SVM model
svm_model = SVC(probability=True, random_state=42)
svm_model.fit(X_train_tfidf, y_train)
# Save the SVM model, TF-IDF vectorizer, and label encoder
joblib.dump(svm_model, svm_model_path)
joblib.dump(vectorizer, vectorizer_path)
joblib.dump(le, label_encoder_path)
print("Models trained and saved to disk.")
# Single-label classification function
def classify_text_svm(text):
text_tfidf = vectorizer.transform([text])
predicted_class_index = svm_model.predict(text_tfidf)[0]
predicted_category = le.inverse_transform([predicted_class_index])[0]
return predicted_category
# Multi-label classification function (returning top N predictions based on probabilities)
def classify_text_svm_multi(text, top_n=3):
text_tfidf = vectorizer.transform([text])
probabilities = svm_model.predict_proba(text_tfidf)[0]
top_n_indices = np.argsort(probabilities)[::-1][:top_n] # Get indices of top N predictions
top_n_categories = le.inverse_transform(top_n_indices)
return top_n_categories
|