KameronB/SITTCIC-roBERTa-f16

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F

# Load the model and tokenizer from Hugging Face
model_name = "KameronB/SITTCIC-roBERTa-f32"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name)

print("Loading model...")
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Set model to evaluation mode
model.eval()

print(f"Model loaded successfully!")
print(f"Model architecture: {model.config.architectures}")
print(f"Number of labels: {model.config.num_labels}")

# Check if CUDA is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"Using device: {device}")




# Function to make predictions
def predict_text(text, model, tokenizer, device):
    """
    Make a prediction on input text using the loaded model
    """
    # Tokenize the input text
    inputs = tokenizer(text, 
                      return_tensors="pt", 
                      truncation=True, 
                      padding=True, 
                      max_length=512)
    
    # Move inputs to the same device as the model
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # Make prediction
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        
        # Apply softmax to get probabilities
        probabilities = F.softmax(logits, dim=-1)
        
        # Get predicted class
        predicted_class = torch.argmax(logits, dim=-1).item()
        confidence = probabilities[0][predicted_class].item()
        
    return predicted_class, confidence, probabilities.cpu().numpy()

# Example inference
sample_texts = [
    "Issue resolved",
    "Cleared the laptop's cache and cookies then restarted it.", 
]

print("Making predictions on sample texts:")
print("-" * 50)

for i, text in enumerate(sample_texts, 1):
    predicted_class, confidence, probabilities = predict_text(text, model, tokenizer, device)
    
    print(f"Text {i}: {text}")
    print(f"Predicted Class: {predicted_class}")
    print(f"Confidence: {confidence:.4f}")
    print(f"All probabilities: {probabilities[0]}")
    print("-" * 50)