Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from transformers import AutoModelForSequenceClassification, DebertaV2Tokenizer
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 CLASSES = {
     'yes': 0,
     'irrelevant': 1,
@@ -12,11 +13,17 @@ model = AutoModelForSequenceClassification.from_pretrained('MrPio/TheSeagullStor
 model.eval()
 story = open('story.txt').read().replace("\n\n", "\n").replace("\n", " ").strip()
 
+if torch.cuda.is_available():
+    model.half()
 
 def ask(question):
-    inputs = tokenizer(story, question, truncation=True, padding=True)
-    prediction = torch.softmax(model(**inputs).logits, 1).squeeze()
-    return {c: round(prediction[i].item(), 3) for c, i in CLASSES}
+  with torch.no_grad():
+    input = tokenizer(story, question, truncation=True, padding=True,return_tensors="pt")
+    input = {key: value.to(device) for key, value in input.items()}
+    output=model(**input)
+    prediction = torch.softmax(output.logits, 1).squeeze()
+    print(prediction)
+    return {c: round(prediction[i].item(), 3) for c, i in CLASSES.items()}
 
 
 gradio = gr.Interface(