| class CustomProvider: | |
| def __init__(self, model, tokenizer): | |
| self.model = model | |
| self.tokenizer = tokenizer | |
| def __call__(self, text): | |
| inputs = self.tokenizer(text, return_tensors="pt") | |
| outputs = self.model(**inputs) | |
| return outputs.logits.argmax(dim=-1).item() | |