justKevv commited on
Commit
a5eabc7
·
1 Parent(s): 73b614e

Refactor model loading: implement lazy loading for models and update .gitignore

Browse files
Files changed (3) hide show
  1. .gitignore +1 -1
  2. app/core/models.py +29 -25
  3. app/services/ranking.py +7 -1
.gitignore CHANGED
@@ -1,2 +1,2 @@
1
  __pycache__
2
- venv
 
1
  __pycache__
2
+ .venv
app/core/models.py CHANGED
@@ -1,14 +1,22 @@
 
 
1
  import joblib
2
  from sentence_transformers import SentenceTransformer
3
  import os
4
  import requests
5
- from huggingface_hub import snapshot_download # <-- Add this import
6
 
7
  TEMP_DIR = "/tmp"
8
 
9
- # This function for your .pkl files is still correct and needed
10
- def download_and_load_model(model_url, model_filename):
11
- # ... (no changes needed in this function)
 
 
 
 
 
 
12
  local_path = os.path.join(TEMP_DIR, model_filename)
13
  if not os.path.exists(local_path):
14
  print(f"Downloading model from {model_url} to {local_path}...")
@@ -28,27 +36,19 @@ def download_and_load_model(model_url, model_filename):
28
  print(f"Failed to load model {local_path}. Error: {e}")
29
  return None
30
 
31
- # --- NEW FUNCTION FOR THE SENTENCE TRANSFORMER ---
32
  def get_sentence_transformer(model_name='all-MiniLM-L6-v2'):
33
- """
34
- Downloads the SentenceTransformer model to /tmp if it doesn't exist,
35
- then loads it from there.
36
- """
37
  local_model_path = os.path.join(TEMP_DIR, model_name)
38
-
39
  if not os.path.exists(local_model_path):
40
  print(f"Downloading SentenceTransformer model '{model_name}' to {local_model_path}...")
41
- # Use snapshot_download to get all files for the model from Hugging Face
42
  try:
43
  snapshot_download(repo_id=f"sentence-transformers/{model_name}",
44
  local_dir=local_model_path,
45
- local_dir_use_symlinks=False) # This is important for Vercel
46
  print("Download complete.")
47
  except Exception as e:
48
  print(f"Failed to download SentenceTransformer model. Error: {e}")
49
  return None
50
-
51
- # Load the model from the local path in /tmp
52
  try:
53
  print(f"Loading SentenceTransformer model from {local_model_path}...")
54
  return SentenceTransformer(local_model_path)
@@ -56,30 +56,34 @@ def get_sentence_transformer(model_name='all-MiniLM-L6-v2'):
56
  print(f"Failed to load SentenceTransformer model from {local_model_path}. Error: {e}")
57
  return None
58
 
59
- # --- Main Model Loading Logic ---
60
- try:
61
- # Your .pkl model loading remains the same
 
 
 
 
 
 
 
62
  TFIDF_URL = "https://pub-4a389a9b2dc842a2a55678d2db0ec0c6.r2.dev/tfidf_vectorizer.pkl"
63
  LE_URL = "https://pub-4a389a9b2dc842a2a55678d2db0ec0c6.r2.dev/label_encoder.pkl"
64
  RF_MODEL_URL = "https://pub-4a389a9b2dc842a2a55678d2db0ec0c6.r2.dev/random_forest_model.pkl"
65
 
66
- tfidf_vectorizer = download_and_load_model(TFIDF_URL, "tfidf_vectorizer.pkl")
67
- le = download_and_load_model(LE_URL, "label_encoder.pkl")
68
- rf_model = download_and_load_model(RF_MODEL_URL, "random_forest_model.pkl")
69
 
70
  if all([tfidf_vectorizer, le, rf_model]):
71
  print("Classification models loaded successfully.")
72
  else:
73
  print("One or more classification models failed to load.")
74
 
75
- # --- THIS IS THE LINE TO CHANGE ---
76
- # Old line: sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
77
- # New line:
78
  sentence_model = get_sentence_transformer()
79
  if sentence_model:
80
  print("SentenceTransformer model loaded successfully.")
81
  else:
82
  print("SentenceTransformer model failed to load.")
83
 
84
- except Exception as e:
85
- print(f"An unexpected error occurred during model loading: {e}")
 
1
+ # app/core/models.py
2
+
3
  import joblib
4
  from sentence_transformers import SentenceTransformer
5
  import os
6
  import requests
7
+ from huggingface_hub import snapshot_download
8
 
9
  TEMP_DIR = "/tmp"
10
 
11
+ # --- Initialize models as None ---
12
+ # They will be loaded into these global variables later.
13
+ tfidf_vectorizer = None
14
+ le = None
15
+ rf_model = None
16
+ sentence_model = None
17
+
18
+ # This function remains the same
19
+ def download_and_load_pkl_model(model_url, model_filename):
20
  local_path = os.path.join(TEMP_DIR, model_filename)
21
  if not os.path.exists(local_path):
22
  print(f"Downloading model from {model_url} to {local_path}...")
 
36
  print(f"Failed to load model {local_path}. Error: {e}")
37
  return None
38
 
39
+ # This function remains the same
40
  def get_sentence_transformer(model_name='all-MiniLM-L6-v2'):
 
 
 
 
41
  local_model_path = os.path.join(TEMP_DIR, model_name)
 
42
  if not os.path.exists(local_model_path):
43
  print(f"Downloading SentenceTransformer model '{model_name}' to {local_model_path}...")
 
44
  try:
45
  snapshot_download(repo_id=f"sentence-transformers/{model_name}",
46
  local_dir=local_model_path,
47
+ local_dir_use_symlinks=False)
48
  print("Download complete.")
49
  except Exception as e:
50
  print(f"Failed to download SentenceTransformer model. Error: {e}")
51
  return None
 
 
52
  try:
53
  print(f"Loading SentenceTransformer model from {local_model_path}...")
54
  return SentenceTransformer(local_model_path)
 
56
  print(f"Failed to load SentenceTransformer model from {local_model_path}. Error: {e}")
57
  return None
58
 
59
+ # --- NEW LAZY LOADING FUNCTION ---
60
+ def load_all_models():
61
+ """
62
+ This function loads all models into the global variables.
63
+ It will only be called when the models are first needed.
64
+ """
65
+ global tfidf_vectorizer, le, rf_model, sentence_model
66
+ print("--- First request received, initiating model loading... ---")
67
+
68
+ # Load the classification models
69
  TFIDF_URL = "https://pub-4a389a9b2dc842a2a55678d2db0ec0c6.r2.dev/tfidf_vectorizer.pkl"
70
  LE_URL = "https://pub-4a389a9b2dc842a2a55678d2db0ec0c6.r2.dev/label_encoder.pkl"
71
  RF_MODEL_URL = "https://pub-4a389a9b2dc842a2a55678d2db0ec0c6.r2.dev/random_forest_model.pkl"
72
 
73
+ tfidf_vectorizer = download_and_load_pkl_model(TFIDF_URL, "tfidf_vectorizer.pkl")
74
+ le = download_and_load_pkl_model(LE_URL, "label_encoder.pkl")
75
+ rf_model = download_and_load_pkl_model(RF_MODEL_URL, "random_forest_model.pkl")
76
 
77
  if all([tfidf_vectorizer, le, rf_model]):
78
  print("Classification models loaded successfully.")
79
  else:
80
  print("One or more classification models failed to load.")
81
 
82
+ # Load the sentence transformer
 
 
83
  sentence_model = get_sentence_transformer()
84
  if sentence_model:
85
  print("SentenceTransformer model loaded successfully.")
86
  else:
87
  print("SentenceTransformer model failed to load.")
88
 
89
+ print("--- Model loading complete. ---")
 
app/services/ranking.py CHANGED
@@ -3,7 +3,7 @@ import json
3
  import os
4
  from geopy.geocoders import Nominatim
5
 
6
- from ..core.models import tfidf_vectorizer, le, rf_model, sentence_model
7
  from ..schemas.recommendation import RecommendationRequest
8
  from ..utils.text import clean_resume
9
 
@@ -55,6 +55,9 @@ def geo_coords(city_name: str) -> tuple | None:
55
 
56
  def get_category_prediction(profile_text: str) -> str:
57
  """Processes text and predicts the job category."""
 
 
 
58
  cleaned_text = profile_text.lower()
59
  vectorized_text = tfidf_vectorizer.transform([cleaned_text])
60
  prediction_encoded = rf_model.predict(vectorized_text)[0]
@@ -64,6 +67,9 @@ def get_category_prediction(profile_text: str) -> str:
64
  def get_ranked_internships(request: RecommendationRequest) -> list[int]:
65
  """Performs two-stage ranking with dynamic geocoding."""
66
 
 
 
 
67
  profile_text_to_encode = request.profile_text
68
 
69
  if request.predicted_category:
 
3
  import os
4
  from geopy.geocoders import Nominatim
5
 
6
+ from ..core.models import load_all_models, tfidf_vectorizer, le, rf_model, sentence_model
7
  from ..schemas.recommendation import RecommendationRequest
8
  from ..utils.text import clean_resume
9
 
 
55
 
56
  def get_category_prediction(profile_text: str) -> str:
57
  """Processes text and predicts the job category."""
58
+ if rf_model is None:
59
+ load_all_models()
60
+
61
  cleaned_text = profile_text.lower()
62
  vectorized_text = tfidf_vectorizer.transform([cleaned_text])
63
  prediction_encoded = rf_model.predict(vectorized_text)[0]
 
67
  def get_ranked_internships(request: RecommendationRequest) -> list[int]:
68
  """Performs two-stage ranking with dynamic geocoding."""
69
 
70
+ if sentence_model is None:
71
+ load_all_models()
72
+
73
  profile_text_to_encode = request.profile_text
74
 
75
  if request.predicted_category: