Update app.py

app.py

@@ -1,6 +1,9 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForSequenceClassification, DebertaV2Tokenizer
+import tensorflow as tf
+from transformers import AutoModelForSequenceClassification, DebertaV2Tokenizer,TFAutoModelForSequenceClassification
+
+USE_TENSORFLOW=True
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 CLASSES = {
@@ -9,21 +12,24 @@ CLASSES = {
     'no': 2,
 }
 tokenizer = DebertaV2Tokenizer.from_pretrained('cross-encoder/nli-deberta-v3-base', do_lower_case=True)
-model = AutoModelForSequenceClassification.from_pretrained('MrPio/TheSeagullStory-nli-deberta-v3-base')
-model.eval()
+model = TFAutoModelForSequenceClassification.from_pretrained('MrPio/TheSeagullStory-nli-deberta-v3-base',dtype=tf.float16) if USE_TENSORFLOW else AutoModelForSequenceClassification.from_pretrained('MrPio/TheSeagullStory-nli-deberta-v3-base')
+if not USE_TENSORFLOW:
+    model.eval()
+    if torch.cuda.is_available():
+        model.half()
 story = open('story.txt').read().replace("\n\n", "\n").replace("\n", " ").strip()
 
-if torch.cuda.is_available():
-    model.half()
-
 def ask(question):
-  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()}
+    input = tokenizer(story, question, truncation=True, padding=True,return_tensors='tf' if USE_TENSORFLOW else 'pt')
+    if not USE_TENSORFLOW:
+        input = {key: value.to(device) for key, value in input.items()}
+        output=model(**input)
+        prediction = torch.softmax(output.logits, 1).squeeze()
+        return {c: round(prediction[i].item(), 3) for c, i in CLASSES.items()}
+    else:
+      output=model(input, training=False)
+      prediction = tf.nn.softmax(output.logits, axis=-1).numpy().squeeze()
+      return {c: round(prediction[i], 3) for c, i in CLASSES.items()}
 
 
 gradio = gr.Interface(