Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3import torch.nn.functional as F
4
5# Load the model and tokenizer from Hugging Face
6model_name = "KameronB/SITTCIC-roBERTa-f32"
7
8print("Loading tokenizer...")
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10
11print("Loading model...")
12model = AutoModelForSequenceClassification.from_pretrained(model_name)
13
14# Set model to evaluation mode
15model.eval()
16
17print(f"Model loaded successfully!")
18print(f"Model architecture: {model.config.architectures}")
19print(f"Number of labels: {model.config.num_labels}")
20
21# Check if CUDA is available
22device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23model.to(device)
24print(f"Using device: {device}")
25
26
27
28
29# Function to make predictions
30def predict_text(text, model, tokenizer, device):
31 """
32 Make a prediction on input text using the loaded model
33 """
34 # Tokenize the input text
35 inputs = tokenizer(text,
36 return_tensors="pt",
37 truncation=True,
38 padding=True,
39 max_length=512)
40
41 # Move inputs to the same device as the model
42 inputs = {k: v.to(device) for k, v in inputs.items()}
43
44 # Make prediction
45 with torch.no_grad():
46 outputs = model(**inputs)
47 logits = outputs.logits
48
49 # Apply softmax to get probabilities
50 probabilities = F.softmax(logits, dim=-1)
51
52 # Get predicted class
53 predicted_class = torch.argmax(logits, dim=-1).item()
54 confidence = probabilities[0][predicted_class].item()
55
56 return predicted_class, confidence, probabilities.cpu().numpy()
57
58# Example inference
59sample_texts = [
60 "Issue resolved",
61 "Cleared the laptop's cache and cookies then restarted it.",
62]
63
64print("Making predictions on sample texts:")
65print("-" * 50)
66
67for i, text in enumerate(sample_texts, 1):
68 predicted_class, confidence, probabilities = predict_text(text, model, tokenizer, device)
69
70 print(f"Text {i}: {text}")
71 print(f"Predicted Class: {predicted_class}")
72 print(f"Confidence: {confidence:.4f}")
73 print(f"All probabilities: {probabilities[0]}")
74 print("-" * 50)