wehub-resource-sync
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
# For full-parameter training, please refer to:
|
|
# https://github.com/modelscope/ms-swift/blob/main/examples/infer/demo_embedding.py
|
|
|
|
import torch
|
|
|
|
from swift.infer_engine import InferRequest, TransformersEngine
|
|
|
|
|
|
def run_qwen3_emb():
|
|
engine = TransformersEngine(
|
|
'Qwen/Qwen3-Embedding-4B',
|
|
task_type='embedding',
|
|
attn_impl='flash_attention_2',
|
|
adapters=['output/vx-xxx/checkpoint-xxx'])
|
|
|
|
infer_requests = [
|
|
InferRequest(messages=[
|
|
{
|
|
'role': 'user',
|
|
'content': 'A dog sleeping under a table.'
|
|
},
|
|
]),
|
|
InferRequest(messages=[
|
|
{
|
|
'role': 'user',
|
|
'content': 'a dog napping under a small table.'
|
|
},
|
|
]),
|
|
InferRequest(messages=[
|
|
{
|
|
'role': 'user',
|
|
'content': 'a cat napping under a small tree.'
|
|
},
|
|
])
|
|
]
|
|
resp_list = engine.infer(infer_requests)
|
|
embedding0 = torch.tensor(resp_list[0].data[0].embedding)
|
|
embedding1 = torch.tensor(resp_list[1].data[0].embedding)
|
|
embedding2 = torch.tensor(resp_list[2].data[0].embedding)
|
|
embedding = torch.stack([embedding0, embedding1, embedding2])
|
|
print(f'scores: {embedding @ embedding.T}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run_qwen3_emb()
|