import gradio as gr
from transformers import (
      BartForConditionalGeneration,
      BartTokenizer
)

model_name = 'unlisboa/bart_qa_assistant'
tokenizer = BartTokenizer.from_pretrained(model_name)
device = get_device()
model = BartForConditionalGeneration.from_pretrained(model_name).to(device)
model.eval()

def example(question, censor):
    print(question, censor)
    return question + str(censor)
    
examples = [["What's the meaning of life?", True]]

checkbox = gr.Checkbox(value=True, label="should censor output")
question_input = gr.Textbox(lines=2, label='Question:')

model_input = tokenizer(question_input, truncation=True, padding=True, return_tensors="pt")
generated_answers_encoded = model.generate(input_ids=model_input["input_ids"].to(device),
                                           attention_mask=model_input["attention_mask"].to(device),
                                           #bad_words_ids=bad_words_ids,
                                           force_words_ids=None,
                                           min_length=1,
                                           max_length=100,
                                           do_sample=True,
                                           early_stopping=True,
                                           num_beams=4,
                                           temperature=1.0,
                                           top_k=None,
                                           top_p=None,
                                           # eos_token_id=tokenizer.eos_token_id,
                                           no_repeat_ngram_size=2,
                                           num_return_sequences=1,
                                           return_dict_in_generate=True,
                                           output_scores=True)

response = tokenizer.batch_decode(generated_answers_encoded['sequences'], skip_special_tokens=True,
                                  clean_up_tokenization_spaces=True)[0]

answer_output = gr.Textbox(lines=2, label='Answer:')
gr.Interface(fn=example, inputs=[question_input, checkbox], outputs=[answer_output], allow_flagging="never", examples=examples).launch()