Update app.py

app.py

@@ -4,6 +4,11 @@ from transformers import (
       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...
@@ -27,6 +32,12 @@ 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),
@@ -36,6 +47,7 @@ def run_bart(question, censor):
                                                                                       min_length=1,
                                                                                       max_length=100,
                                                                                       do_sample=True,
+                                                                                      bad_words_ids=bad_words_ids
                                                                                       early_stopping=True,
                                                                                       num_beams=4,
                                                                                       temperature=1.0,