Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import ( | |
| BartForConditionalGeneration, | |
| BartTokenizer | |
| ) | |
| import torch | |
| import json | |
| def read_json_file_2_dict(filename, store_dir='.'): | |
| with open(f'{store_dir}/{filename}', 'r', encoding='utf-8') as file: | |
| return json.load(file) | |
| def get_device(): | |
| # If there's a GPU available... | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| n_gpus = torch.cuda.device_count() | |
| first_gpu = torch.cuda.get_device_name(0) | |
| print(f'There are {n_gpus} GPU(s) available.') | |
| print(f'GPU gonna be used: {first_gpu}') | |
| else: | |
| print('No GPU available, using the CPU instead.') | |
| device = torch.device("cpu") | |
| return device | |
| 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 run_bart(question, censor): | |
| print(question, censor) | |
| if censor: | |
| bad_words = read_json_file_2_dict('bad_words_file.json') | |
| bad_words_ids = tokenizer(bad_words, add_prefix_space=True, add_special_tokens=False).get('input_ids') | |
| else: | |
| bad_words_ids = None | |
| model_input = tokenizer(question, 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, | |
| bad_words_ids=bad_words_ids, | |
| 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] | |
| return response | |
| 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:') | |
| answer_output = gr.Textbox(lines=2, label='Answer:') | |
| gr.Interface(fn=run_bart, inputs=[question_input, checkbox], outputs=[answer_output], allow_flagging="never", examples=examples).launch() |