Update README.md

README.md

@@ -47,6 +47,7 @@ import torch
 import transformers
 from PIL import Image
 import sys
+
 sys.path.append("..")
 from otter.modeling_otter import OtterForConditionalGeneration
 
@@ -130,6 +131,7 @@ def get_response(input_data, prompt: str, model=None, image_processor=None, tens
         return_tensors="pt",
     )
 
+    bad_words_id = model.text_tokenizer(["User:", "GPT1:", "GFT:", "GPT:"], add_special_tokens=False).input_ids
     generated_text = model.generate(
         vision_x=vision_x.to(model.device, dtype=tensor_dtype),
         lang_x=lang_x["input_ids"].to(model.device),
@@ -137,6 +139,7 @@ def get_response(input_data, prompt: str, model=None, image_processor=None, tens
         max_new_tokens=512,
         num_beams=3,
         no_repeat_ngram_size=3,
+        bad_words_ids=bad_words_id,
     )
     parsed_output = (
         model.text_tokenizer.decode(generated_text[0])
@@ -151,6 +154,7 @@ def get_response(input_data, prompt: str, model=None, image_processor=None, tens
     )
     return parsed_output
 
+
 # ------------------- Main Function -------------------
 load_bit = "fp16"
 if load_bit == "fp16":
@@ -184,4 +188,5 @@ while True:
 
     if prompts_input.lower() == "quit":
         break
+
 ```