Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3import numpy as np
4
5query = "Jak dożyć 100 lat?"
6answers = [
7 "Trzeba zdrowo się odżywiać i uprawiać sport.",
8 "Trzeba pić alkohol, imprezować i jeździć szybkimi autami.",
9 "Gdy trwała kampania politycy zapewniali, że rozprawią się z zakazem niedzielnego handlu."
10]
11
12model_name = "sdadas/polish-reranker-roberta-v2"
13tokenizer = AutoTokenizer.from_pretrained(model_name)
14model = AutoModelForSequenceClassification.from_pretrained(
15 model_name,
16 dtype=torch.bfloat16,
17 device_map="cuda"
18)
19texts = [f"{query}</s></s>{answer}" for answer in answers]
20tokens = tokenizer(texts, padding="longest", max_length=512, truncation=True, return_tensors="pt").to("cuda")
21output = model(**tokens)
22results = output.logits.detach().cpu().float().numpy()
23results = np.squeeze(results)
24print(results.tolist())1@article{dadas2024assessing,
2 title={Assessing generalization capability of text ranking models in Polish},
3 author={Sławomir Dadas and Małgorzata Grębowiec},
4 year={2024},
5 eprint={2402.14318},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL}
8}