Spaces:
Sleeping
Sleeping
| 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) | |