Views
No views yet
1from transformers import AutoModelForSequenceClassification, AutoTokenizer
2
3# Set device to CPU or GPU
4import torch
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6
7# Load model and tokenizer
8model_name = "ariannap22/collectiveaction_roberta_simplified_synthetic_weights"
9model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11
12# Define the text you want to predict
13texts = [
14 "We need to stand together for our rights!",
15 "I volunteer at the local food bank.",
16 "This is just a casual discussion."
17]
18
19# Tokenize the input text
20inputs = tokenizer(
21 texts,
22 padding=True, # Pad to the longest sequence in the batch
23 truncation=True, # Truncate sequences longer than the model's max length
24 max_length=512, # Adjust max length as needed
25 return_tensors="pt" # Return PyTorch tensors
26).to(device)
27
28# Perform prediction
29with torch.no_grad():
30 outputs = model(**inputs)
31 logits = outputs.logits # Raw model outputs before softmax
32
33# Convert logits to probabilities (optional)
34probs = torch.nn.functional.softmax(logits, dim=-1)
35
36# Get predicted class indices
37predicted_class_indices = torch.argmax(probs, dim=-1)
38
39# Print results
40for text, idx, prob in zip(texts, predicted_class_indices, probs):
41 print(f"Text: {text}")
42 print(f"Predicted Class Index: {idx.item()}")
43 print(f"Probabilities: {prob.tolist()}")
44 print("---")
45