Spaces:
Runtime error
Runtime error
| 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() |