resume-api / app /core /models.py
justKevv's picture
Add application file
5e0b388
Raw
History Blame Contribute Delete
1.22 kB
import joblib
from sentence_transformers import SentenceTransformer
import os
MODEL_DIR = "models"
# Initialize variables
tfidf_vectorizer = None
le = None
rf_model = None
sentence_model = None
try:
tfidf_vectorizer = joblib.load(os.path.join(MODEL_DIR, "tfidf_vectorizer.pkl"))
le = joblib.load(os.path.join(MODEL_DIR, "label_encoder.pkl"))
rf_model = joblib.load(os.path.join(MODEL_DIR, "random_forest_model.pkl"))
print("Classification models loaded.")
# Set cache directory to a writable location for Hugging Face Spaces
os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
os.environ['HF_HOME'] = '/tmp/hf_home'
sentence_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
print("SentenceTransformer model loaded.")
except FileNotFoundError as e:
print(f"MODEL LOADING ERROR: {e}")
print("Make sure the .pkl files are in the 'models' directory.")
raise e # Re-raise to prevent the application from starting with None models
except Exception as e:
print(f"An unexpected error occurred during model loading: {e}")
raise e # Re-raise to prevent the application from starting with None models