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