api-recruitment-demo / embeddings.py
jonathanjordan21's picture
Update embeddings.py
41220ca verified
# Requires transformers>=4.51.0
import torch
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")
def rank_jobs(job_description, resumes):
task = "Given a resume, retrieve relevant job description that is suitable for the resume"
queries = resumes
documents = job_description
print("[QUERIES]", queries)
print("[DOCUMENTS]", documents)
query_embeddings = model.encode(queries, prompt=task)
document_embeddings = model.encode(documents)
similarity = model.similarity(query_embeddings, document_embeddings)
return documents, similarity[0].tolist()
def rank_jobs_multi(job_description, resumes):
task = "Given a resume, retrieve relevant job description that is suitable for the resume"
queries = resumes
documents = job_description
print("[QUERIES]", queries)
print("[DOCUMENTS]", documents)
query_embeddings = model.encode(queries, prompt=task)
document_embeddings = model.encode(documents)
similarity = model.similarity(query_embeddings, document_embeddings)
return documents, similarity[:,0].tolist()