wehub-resource-sync
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import torch
|
|
|
|
from swift.infer_engine import InferRequest, TransformersEngine
|
|
|
|
|
|
def run_qwen3_reranker():
|
|
engine = TransformersEngine(
|
|
'Qwen/Qwen3-Reranker-4B',
|
|
task_type='generative_reranker',
|
|
torch_dtype=torch.float16,
|
|
attn_impl='flash_attention_2')
|
|
|
|
infer_request = InferRequest(
|
|
messages=[{
|
|
'role': 'system',
|
|
'content': 'Given a web search query, retrieve relevant passages that answer the query'
|
|
}, {
|
|
'role': 'user',
|
|
'content': 'What is the capital of China?'
|
|
}, {
|
|
'role': 'assistant',
|
|
'content': 'The capital of China is Beijing.'
|
|
}])
|
|
|
|
response = engine.infer([infer_request])[0]
|
|
print(f'scores: {response.choices[0].message.content}')
|
|
|
|
|
|
def run_qwen3_vl_reranker():
|
|
engine = TransformersEngine(
|
|
'Qwen/Qwen3-VL-Reranker-2B', task_type='generative_reranker', attn_impl='flash_attention_2')
|
|
|
|
infer_request = InferRequest(
|
|
messages=[{
|
|
'role': 'system',
|
|
'content': "Retrieval relevant image or text with user's query"
|
|
}, {
|
|
'role': 'user',
|
|
'content': 'A woman playing with her dog on a beach at sunset.'
|
|
}, {
|
|
'role':
|
|
'assistant',
|
|
'content':
|
|
'<image>A woman shares a joyful moment with her golden retriever on a sun-drenched beach '
|
|
'at sunset, as the dog offers its paw in a heartwarming display of companionship and trust.'
|
|
}],
|
|
images=['https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg'])
|
|
|
|
response = engine.infer([infer_request])[0]
|
|
print(f'scores: {response.choices[0].message.content}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# run_qwen3_reranker()
|
|
run_qwen3_vl_reranker()
|