Update app.py

app.py

@@ -1,31 +1,29 @@
-import torch
-from transformers import BigBirdForQuestionAnswering, BigBirdTokenizerFast
+import jax.numpy as jnp
+from transformers import FlaxBigBirdForQuestionAnswering, BigBirdTokenizerFast
 import gradio as gr
 
 FLAX_MODEL_ID = "vasudevgupta/flax-bigbird-natural-questions"
 
-if __name__ == "__main__":
-  device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
-  
-  model = BigBirdForQuestionAnswering.from_pretrained(FLAX_MODEL_ID, block_size=64, num_random_blocks=3, from_flax=True)
-  model.to(device)
+if __name__ == "__main__":  
+  model = FlaxBigBirdForQuestionAnswering.from_pretrained(FLAX_MODEL_ID, block_size=64, num_random_blocks=3)
   tokenizer = BigBirdTokenizerFast.from_pretrained(FLAX_MODEL_ID)
   
+  @jax.jit
+  def forward(*args, **kwargs):
+    return model(*args, **kwargs)
+  
   def get_answer(question, context):
-    encoding = tokenizer(question, context, return_tensors="pt", max_length=4096, padding="max_length", truncation=True)
-    input_ids = encoding.input_ids.to(device)
-    attention_mask = encoding.attention_mask.to(device)
-
-    with torch.no_grad():
-        start_scores, end_scores = model(input_ids=input_ids, attention_mask=attention_mask).to_tuple()
-
-    # Let's take the most likely token using `argmax` and retrieve the answer
-    all_tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"].squeeze().tolist())
-
-    answer_tokens = all_tokens[torch.argmax(start_scores): torch.argmax(end_scores)+1]
-    answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
-
-    return answer
+  
+      encoding = tokenizer(question, context, return_tensors="jax", max_length=4096, padding="max_length", truncation=True)
+      start_scores, end_scores = forward(**encoding).to_tuple()
+  
+      # Let's take the most likely token using `argmax` and retrieve the answer
+      all_tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"][0].tolist())
+  
+      answer_tokens = all_tokens[jnp.argmax(start_scores): jnp.argmax(end_scores)+1]
+      answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
+  
+      return answer
 
     default_context = "BigBird Pegasus just landed! Thanks to Vasudev Gupta, BigBird Pegasus from Google AI is merged into HuggingFace Transformers. Check it out today!!!"
     question = gr.inputs.TextBox(lines=2, default="Who added BigBird to HuggingFace Transformers?", label="Question")