chore: add bad words

app.py

@@ -7,6 +7,7 @@ from transformers import (
     AutoModelForPreTraining,
     AutoProcessor,
     AutoConfig,
+    PreTrainedTokenizerFast
 )
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
@@ -18,10 +19,16 @@ assert MODEL_NAME is not None
 MODEL_PATH = hf_hub_download(repo_id=MODEL_NAME, filename="model.safetensors")
 DEVICE = torch.device("cuda")
 
+BAD_WORD_KEYWORDS = ["(medium)"]
 
 def fix_compiled_state_dict(state_dict: dict):
     return {k.replace("._orig_mod.", "."): v for k, v in state_dict.items()}
 
+def get_bad_words_ids(tokenizer: PreTrainedTokenizerFast):
+    ids = [
+        [id] for token, id in tokenizer.vocab.items() if any(word in token for BAD_WORD_KEYWORDS)
+    ]
+    return ids
 
 def prepare_models():
     config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
@@ -44,6 +51,7 @@ def prepare_models():
 
 def demo():
     model, processor = prepare_models()
+    ban_ids = get_bad_words_ids(processor.decoder_tokenizer)
 
     @spaces.GPU(duration=5)
     @torch.inference_mode()
@@ -83,6 +91,7 @@ def demo():
             top_p=top_p,
             eos_token_id=processor.decoder_tokenizer.eos_token_id,
             pad_token_id=processor.decoder_tokenizer.pad_token_id,
+            bad_words_ids=ban_ids,
         )
         elapsed = time.time() - start_time