rwcuffney commited on
Commit
c51d635
1 Parent(s): 6f0acf6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_dataset
2
+
3
+ dataset = load_dataset('rwcuffney/pick_a_card_test', batch_size=32, shuffle=True)
4
+
5
+ from transformers import AutoModelForSequenceClassification
6
+
7
+ model = AutoModelForSequenceClassification.from_pretrained('rwcuffney/autotrain-pick_a_card-3726099224')
8
+
9
+ from transformers import AutoTokenizer
10
+
11
+ tokenizer = AutoTokenizer.from_pretrained('rwcuffney/autotrain-pick_a_card-3726099224')
12
+
13
+ def preprocess_text(text):
14
+ encoded = tokenizer(text, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
15
+ return encoded
16
+
17
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
18
+ model.to(device)
19
+ model.eval()
20
+
21
+ for batch in dataset:
22
+ # Preprocess the text
23
+ text = batch['text']
24
+ inputs = preprocess_text(text)
25
+ inputs = inputs.to(device)
26
+
27
+ # Make predictions
28
+ with torch.no_grad():
29
+ outputs = model(**inputs)
30
+ predicted_classes = torch.argmax(outputs.logits, dim=-1)
31
+
32
+ # Print the predicted class labels
33
+ predicted_labels = [dataset.features['label'].names[i] for i in predicted_classes]
34
+ print(predicted_labels)
35
+
36
+
37
+