Views
No views yet
1from peft import PeftModel
2from transformers import AutoModelForSequenceClassification, AutoTokenizer
3import torch
4import torch.nn.functional as F
5
6label_map = {0: "clean sentence", 1: "out of dictionary words", 2: "word salad", 3: "number gibberish", 4: "mild gibberish"}
7
8def intlabel_to_strlabel(label):
9 return label_map[label]
10
11def get_prediction(model, tokenzier, input, label_fn):
12 infer_inputs = tokenzier(input, return_tensors="pt")
13
14 infer_device = model.device
15 infer_inputs = {key: value.to(infer_device) for key, value in infer_inputs.items()}
16
17 with torch.no_grad():
18 outputs = model(**infer_inputs)
19
20 logits = outputs.logits
21
22 probabilities = torch.nn.functional.softmax(logits, dim=-1)
23 predicted_class = torch.argmax(probabilities, dim=-1).item()
24
25 probabilities = F.softmax(logits, dim=-1)
26 predicted_index = torch.argmax(probabilities, dim=1).item()
27 predicted_prob = probabilities[0][predicted_index].item()
28
29 label = label_fn(predicted_class)
30
31 output = {"label": label, "score": predicted_prob}
32
33 return output
34
35lora_weights = "sunnysingh1011/gibberish-detection"
36tokenizer_path = "sunnysingh1011/gibberish-detection"
37base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=len(label_map))
38
39device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
40
41inference_model = PeftModel.from_pretrained(base_model, lora_weights).to(device)
42inference_tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
43
44input = "Jumped quick dog over lazy the fox brown."
45get_prediction(inference_model, inference_tokenizer, input, intlabel_to_strlabel)