app.py

import jax
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__":  
  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="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")
  context = gr.inputs.Textbox(lines=10, default=default_context, label="Context")

  gr.Interface(fn=get_answer, inputs=[question, context], outputs="text").launch()