Christian2903 commited on
Commit
8b2ba20
·
1 Parent(s): 034476c

Delete handler.py

Browse files
Files changed (1) hide show
  1. handler.py +0 -31
handler.py DELETED
@@ -1,31 +0,0 @@
1
- from transformers import AutoTokenizer, AutoModelForSequenceClassification
2
- from transformers import DataCollatorWithPadding
3
- from torch.nn.functional import softmax
4
-
5
- import torch
6
- from typing import Any, Dict, List
7
-
8
-
9
- class EndpointHandler:
10
- def __init__(self, path=""):
11
- self.tokenizer = AutoTokenizer.from_pretrained(path)
12
- self.model = AutoModelForSequenceClassification.from_pretrained(path)
13
- self.tokenizer.pad_token = self.tokenizer.eos_token
14
- self.model.config.pad_token_id = self.tokenizer.pad_token_id
15
-
16
- def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
17
- batch_of_strings = data["inputs"]
18
-
19
- tokens = self.tokenizer(
20
- batch_of_strings, padding=True, truncation=True, return_tensors="pt"
21
- )
22
-
23
- # Calculate the loss
24
- with torch.no_grad():
25
- outputs = self.model(**tokens)
26
-
27
- probabilities = softmax(outputs.logits, dim=1)
28
-
29
- return {
30
- "predictions": [pred[0] for pred in probabilities.tolist()],
31
- }