llmengine/qwen3_reranker.py
2025-06-19 17:23:28 +08:00

11 lines
412 B
Python

import torch
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
from llmengine.base_reranker import BaseReranker
class Qwen3Reranker(BaseReranker):
def __init__(self, model_id, max_length=8096):
self.odel_id = model_id
self.tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side='left')
self.model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-Reranker-0.6B").eval()