Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4del_symbs = ["?","!",".",","]
5classes = ["dialog","trouble","question","about_user","about_model","instruct"]
6
7device = torch.device("cuda")
8model_name = 'TeraSpace/replica_classification'
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels = len(classes)).to(device)
11
12while True:
13 text = input("=>").lower()
14 for del_symb in del_symbs:
15 text = text.replace(del_symb,"")
16
17 inputs = tokenizer(text, truncation=True, max_length=512, padding='max_length',
18 return_tensors='pt').to(device)
19 with torch.no_grad():
20 logits = model(**inputs).logits
21 probas = list(torch.sigmoid(logits)[0].cpu().detach().numpy())
22
23 out = classes[probas.index(max(probas))]
24 print(out)