Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4model_id = "cnmoro/BertTiny-Reranker-EnPt"
5model = AutoModelForSequenceClassification.from_pretrained(
6 model_id,
7 num_labels=2
8)
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11model.to(device)
12
13template = "Query: {query}\nSentence: {document}"
14
15def rank(query, documents, normalize_scores=True):
16 texts = [template.format(query=query, document=document) for document in documents]
17
18 inputs = tokenizer(
19 texts,
20 add_special_tokens=True,
21 max_length=512,
22 truncation=True,
23 padding=True,
24 return_tensors="pt",
25 )
26
27 input_ids = inputs["input_ids"].to(device)
28 attention_mask = inputs["attention_mask"].to(device)
29
30 model.eval()
31 with torch.no_grad():
32 outputs = model(input_ids, attention_mask=attention_mask)
33 logits = outputs.logits
34 probabilities = torch.softmax(logits, dim=1)
35
36 # Get the predicted classes and confidence scores
37 predicted_classes = torch.argmax(probabilities, dim=1).tolist()
38 confidences = probabilities.max(dim=1).values.tolist()
39
40 # Construct the results
41 results = [
42 {"prediction": pred, "confidence": conf}
43 for pred, conf in zip(predicted_classes, confidences)
44 ]
45
46 final_results = []
47 for document, result in zip(documents, results):
48 # If the prediction is 0, then get the score as 1 - confidence
49 if result['prediction'] == 0:
50 result['confidence'] = 1 - result['confidence']
51 final_results.append((document, result['confidence']))
52
53 # Sort by the confidence score, descending
54 sorted_results = sorted(final_results, key=lambda x: x[1], reverse=True)
55
56 if normalize_scores:
57 total_score = sum([result[1] for result in sorted_results])
58 sorted_results = [(result[0], result[1] / total_score) for result in sorted_results]
59
60 return sorted_results
61
62# Sample - 1
63query = "O que é o Pantanal?"
64documents = [
65 "É um dos ecossistemas mais ricos em biodiversidade do mundo, abrigando uma grande variedade de espécies animais e vegetais.",
66 "Sua beleza natural, com rios e lagos interligados, atrai turistas de todo o mundo.",
67 "O Pantanal sofre com impactos ambientais, como a exploração mineral e o desmatamento.",
68 "O Pantanal é uma extensa planície alagável localizada na América do Sul, principalmente no Brasil, mas também em partes da Bolívia e Paraguai.",
69 "É um local com importância histórica e cultural para as populações locais.",
70 "O Pantanal é um importante habitat para diversas espécies de animais, inclusive aves migratórias."
71]
72rank(query, documents)
73# [('O Pantanal é uma extensa planície alagável localizada na América do Sul, principalmente no Brasil, mas também em partes da Bolívia e Paraguai.',
74# 0.34217479171361365),
75# ('O Pantanal é um importante habitat para diversas espécies de animais, inclusive aves migratórias.',
76# 0.3310142292228493),
77# ('O Pantanal sofre com impactos ambientais, como a exploração mineral e o desmatamento.',
78# 0.2649093801652631),
79# ('É um local com importância histórica e cultural para as populações locais.',
80# 0.027501075607910826),
81# ('É um dos ecossistemas mais ricos em biodiversidade do mundo, abrigando uma grande variedade de espécies animais e vegetais.',
82# 0.02122344629201432),
83# ('Sua beleza natural, com rios e lagos interligados, atrai turistas de todo o mundo.',
84# 0.013177076998348802)]
85
86# Sample - 2
87query = "What is the speed of light?"
88documents = [
89 "Isaac Newton's laws of motion and gravity laid the groundwork for classical mechanics.",
90 "The theory of relativity, proposed by Albert Einstein, has revolutionized our understanding of space, time, and gravity.",
91 "The Earth orbits the Sun at an average distance of about 93 million miles, taking roughly 365.25 days to complete one revolution.",
92 "The speed of light in a vacuum is approximately 299,792 kilometers per second (km/s), or about 186,282 miles per second.",
93 "Light can be described as both a wave and a particle, a concept known as wave-particle duality."
94]
95rank(query, documents)
96# [('The speed of light in a vacuum is approximately 299,792 kilometers per second (km/s), or about 186,282 miles per second.',
97# 0.2666552015232382),
98# ('The Earth orbits the Sun at an average distance of about 93 million miles, taking roughly 365.25 days to complete one revolution.',
99# 0.25073310834800405),
100# ('Light can be described as both a wave and a particle, a concept known as wave-particle duality.',
101# 0.23357900324953587),
102# ('The theory of relativity, proposed by Albert Einstein, has revolutionized our understanding of space, time, and gravity.',
103# 0.1280649276771122),
104# ("Isaac Newton's laws of motion and gravity laid the groundwork for classical mechanics.",
105# 0.12096775920210973)]