import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer MODEL = "EverMind-AI/skillrep-reranker-0.6b"tok = AutoTokenizer.from_pretrained(MODEL, padding_side="left")model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).cuda().eval() PREFIX = ('<|im_start|>system\nJudge whether the Document meets the requirements ' 'based on the Query and the Instruct provided. Note that the answer can ' 'only be "yes" or "no".<|im_end|>\n<|im_start|>user\n')SUFFIX = '<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n'INSTRUCT = ("Given a coding task description, judge whether the skill document " "is relevant and useful for completing the task") YES, NO = tok.convert_tokens_to_ids("yes"), tok.convert_tokens_to_ids("no")PRE = tok.encode(PREFIX, add_special_tokens=False)SUF = tok.encode(SUFFIX, add_special_tokens=False) def pair(task, skill): doc = f"{skill['name']} | {skill['description'][:500]} | {skill['body'][:2000]}" return f"<Instruct>: {INSTRUCT}\n\n<Query>: {task}\n\n<Document>: {doc}" def score(task, skills, max_length=4096): """One score per skill; higher means more relevant.""" budget = max_length - len(PRE) - len(SUF) # truncate the pair, never the template ids = [PRE + tok.encode(pair(task, s), add_special_tokens=False, truncation=True, max_length=budget) + SUF for s in skills] batch = tok.pad({"input_ids": ids}, return_tensors="pt").to(model.device) with torch.no_grad(): logits = model(**batch).logits[:, -1, :] return (logits[:, YES] - logits[:, NO]).float().tolist() task = "resolve conflicts after a git merge"skills = [ {"name": "resolve-conflicts", "description": "Resolve git merge conflicts.", "body": "# Resolving merge conflicts\n1. Run `git status` to list conflicted files ..."}, {"name": "sourdough", "description": "Bake sourdough bread.", "body": "# Sourdough\nMix flour, water and starter ..."},]print(score(task, skills))# -> [3.0, -4.06]