Spaces:
Sleeping
Sleeping
File size: 1,570 Bytes
6b4436c 565b974 6b4436c 565b974 6b4436c 565b974 6b4436c 565b974 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import evaluate
import datasets
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
class Classicier(evaluate.Measurement):
def _info(self):
return evaluate.MeasurementInfo(
description="",
citation="",
inputs_description="",
features=datasets.Features(
{
"texts": datasets.Value("string", id="sequence"),
}
),
reference_urls=[],
)
def _download_and_prepare(self, dl_manager, device=None):
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the tokenizer and model from the specified repository
self.tokenizer = AutoTokenizer.from_pretrained("AbdulmohsenA/classicier")
self.model = AutoModelForSequenceClassification.from_pretrained("AbdulmohsenA/classicier")
self.model.to(device)
self.device = device
def _compute(self, texts, temperature=2):
device = self.device
inputs = self.tokenizer(
texts,
return_tensors="pt",
truncation=True,
padding='max_length',
max_length=128
).to(device)
with torch.no_grad():
output = self.model(**inputs)
prediction = torch.softmax(output.logits / temperature, dim=-1)
classical_prob = prediction[:, 1].detach().cpu().numpy()
return {"classical_score": classical_prob}
|