Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4# Load model and tokenizer
5tokenizer = AutoTokenizer.from_pretrained("rajinikarcg/RequirementClassifier")
6model = AutoModelForSequenceClassification.from_pretrained("rajinikarcg/RequirementClassifier")
7
8# Prepare input
9text = "The system shall respond within 2 seconds"
10inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
11
12# Get prediction
13with torch.no_grad():
14 outputs = model(**inputs)
15 logits = outputs.logits
16 prediction = torch.argmax(logits, dim=-1).item()
17
18# Map to label
19labels = ["non-requirement", "requirement"]
20print(f"Prediction: {labels[prediction]}")