Spaces:
Sleeping
Sleeping
File size: 1,085 Bytes
bb723d2 |
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 |
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import torch.nn.functional as F
modelName = "negi2725/LegalBertNew"
model = AutoModelForSequenceClassification.from_pretrained(modelName)
tokenizer = AutoTokenizer.from_pretrained(modelName)
model.eval()
def predictVerdict(text: str) -> str:
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding=True)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = F.softmax(logits, dim=-1)
predictedClass = torch.argmax(probabilities, dim=-1).item()
return "guilty" if predictedClass == 1 else "not guilty"
def getConfidence(text: str) -> float:
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding=True)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = F.softmax(logits, dim=-1)
confidence = torch.max(probabilities).item()
return round(confidence, 4)
|