# Requires transformers>=4.51.0
import torch
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")

# queries = "hey"
# documents = [
#     "The capital of China is Beijing.",
#     "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
# ]

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()