webcodelab's picture
Update app/ranker.py
4a365d5 verified
Raw
History Blame Contribute Delete
1.19 kB
# utils/ranker.py
import os
from sentence_transformers import SentenceTransformer, util
from .utils import extract_text_from_pdf
import PyPDF2
# Set safe cache directories for Hugging Face models (especially for read-only environments like Hugging Face Spaces)
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp/hf"
model = SentenceTransformer('all-MiniLM-L6-v2')
def extract_text_from_pdf(file_storage):
reader = PyPDF2.PdfReader(file_storage)
return " ".join([page.extract_text() or '' for page in reader.pages])
def rank_resumes_by_semantic_similarity(resume_files, job_description):
job_embedding = model.encode(job_description, convert_to_tensor=True)
results = []
for resume_file in resume_files:
text = extract_text_from_pdf(resume_file)
resume_embedding = model.encode(text, convert_to_tensor=True)
score = util.cos_sim(job_embedding, resume_embedding).item()
results.append({
"filename": resume_file.filename,
"score": round(score, 4)
})
results.sort(key=lambda x: x['score'], reverse=True)
return results