feat: change for sft version

app.py

@@ -10,7 +10,7 @@ import gradio as gr
 MODEL_NAME = (
     os.environ.get("MODEL_NAME")
     if os.environ.get("MODEL_NAME") is not None
-    else "p1atdev/dart-test-1"
+    else "p1atdev/dart-test-3-sft-1"
 )
 HF_READ_TOKEN = os.environ.get("HF_READ_TOKEN")
 
@@ -43,6 +43,8 @@ CHARACTER_EOS = "</character>"
 GENERAL_BOS = "<general>"
 GENERAL_EOS = "</general>"
 
+INPUT_END = "<|input_end|>"
+
 RATING_BOS_ID = tokenizer.convert_tokens_to_ids(RATING_BOS)
 RATING_EOS_ID = tokenizer.convert_tokens_to_ids(RATING_EOS)
 COPYRIGHT_BOS_ID = tokenizer.convert_tokens_to_ids(COPYRIGHT_BOS)
@@ -52,6 +54,9 @@ CHARACTER_EOS_ID = tokenizer.convert_tokens_to_ids(CHARACTER_EOS)
 GENERAL_BOS_ID = tokenizer.convert_tokens_to_ids(GENERAL_BOS)
 GENERAL_EOS_ID = tokenizer.convert_tokens_to_ids(GENERAL_EOS)
 
+INPUT_END_ID = tokenizer.convert_tokens_to_ids(INPUT_END)
+
+
 assert isinstance(RATING_BOS_ID, int)
 assert isinstance(RATING_EOS_ID, int)
 assert isinstance(COPYRIGHT_BOS_ID, int)
@@ -60,6 +65,7 @@ assert isinstance(CHARACTER_BOS_ID, int)
 assert isinstance(CHARACTER_EOS_ID, int)
 assert isinstance(GENERAL_BOS_ID, int)
 assert isinstance(GENERAL_EOS_ID, int)
+assert isinstance(INPUT_END_ID, int)
 
 SPECIAL_TAGS = [
     BOS,
@@ -72,11 +78,12 @@ SPECIAL_TAGS = [
     CHARACTER_EOS,
     GENERAL_BOS,
     GENERAL_EOS,
+    INPUT_END,
 ]
 
 SPECIAL_TAG_IDS = tokenizer.convert_tokens_to_ids(SPECIAL_TAGS)
 assert isinstance(SPECIAL_TAG_IDS, list)
-
+assert all([token_id != tokenizer.unk_token_id for token_id in SPECIAL_TAG_IDS])
 
 RATING_TAGS = {
     "sfw": "rating:sfw",
@@ -128,29 +135,42 @@ def compose_prompt(
             CHARACTER_EOS,
             GENERAL_BOS,
             general,
+            INPUT_END,
         ]
     )
 
 
 @torch.no_grad()
 def generate(
-    input_text,
-    max_new_tokens=128,
+    input_text: str,
+    max_new_tokens: int = 128,
+    min_new_tokens: int = 0,
     do_sample: bool = True,
     temperature: float = 1.0,
     top_p: float = 1,
     top_k: int = 20,
     num_beams: int = 1,
     bad_words_ids: list[int] | None = None,
+    cfg_scale: float = 1.5,
+    negative_input_text: str | None = None,
 ) -> list[int]:
     inputs = tokenizer(
         input_text,
         return_tensors="pt",
     ).input_ids
+    negative_inputs = (
+        tokenizer(
+            negative_input_text,
+            return_tensors="pt",
+        ).input_ids
+        if negative_input_text is not None
+        else None
+    )
 
     generated = model.generate(
         inputs,
         max_new_tokens=max_new_tokens,
+        min_new_tokens=min_new_tokens,
         do_sample=do_sample,
         temperature=temperature,
         top_p=top_p,
@@ -159,6 +179,8 @@ def generate(
         bad_words_ids=(
             [[token] for token in bad_words_ids] if bad_words_ids is not None else None
         ),
+        negative_prompt_ids=negative_inputs,
+        guidance_scale=cfg_scale,
         no_repeat_ngram_size=1,
     )[0]
 
@@ -171,7 +193,10 @@ def decode_normal(token_ids: list[int], skip_special_tokens: bool = True):
 
 def decode_general_only(token_ids: list[int]):
     token_ids = token_ids[token_ids.index(GENERAL_BOS_ID) :]
-    return tokenizer.decode(token_ids, skip_special_tokens=True)
+    decoded = tokenizer.decode(token_ids, skip_special_tokens=True)
+    tags = [tag for tag in decoded.split(", ")]
+    tags = sorted(tags)
+    return ", ".join(tags)
 
 
 def split_people_tokens_part(token_ids: list[int]):
@@ -242,7 +267,11 @@ def handle_inputs(
     character_tags_list: list[str],
     general_tags: str,
     ban_tags: str,
+    do_cfg: bool = False,
+    cfg_scale: float = 1.5,
+    negative_tags: str = "",
     max_new_tokens: int = 128,
+    min_new_tokens: int = 0,
     temperature: float = 1.0,
     top_p: float = 1.0,
     top_k: int = 20,
@@ -272,17 +301,29 @@ def handle_inputs(
         general=general_tags,
     )
 
-    bad_words_ids = tokenizer.encode_plus(ban_tags).input_ids
+    negative_prompt = compose_prompt(
+        rating=prepare_rating_tags(rating_tags),
+        copyright="",
+        character="",
+        general=negative_tags,
+    )
+
+    bad_words_ids = tokenizer.encode_plus(
+        ban_tags if negative_tags.strip() == "" else ban_tags + ", " + negative_tags
+    ).input_ids
 
     generated_ids = generate(
         prompt,
         max_new_tokens=max_new_tokens,
+        min_new_tokens=min_new_tokens,
         do_sample=True,
         temperature=temperature,
         top_p=top_p,
         top_k=top_k,
         num_beams=num_beams,
         bad_words_ids=bad_words_ids if len(bad_words_ids) > 0 else None,
+        cfg_scale=cfg_scale,
+        negative_input_text=negative_prompt if do_cfg else None,
     )
 
     decoded_normal = decode_normal(generated_ids, skip_special_tokens=True)
@@ -334,7 +375,7 @@ def demo():
                         )
                         copyright_tags_dropdown = gr.Dropdown(
                             label="Copyright tags",
-                            choices=COPYRIGHT_TAGS_LIST,
+                            choices=COPYRIGHT_TAGS_LIST,  # type: ignore
                             value=[],
                             multiselect=True,
                             visible=False,
@@ -363,7 +404,7 @@ def demo():
                         )
                         character_tags_dropdown = gr.Dropdown(
                             label="Character tags",
-                            choices=CHARACTER_TAGS_LIST,
+                            choices=CHARACTER_TAGS_LIST,  # type: ignore
                             value=[],
                             multiselect=True,
                             visible=False,
@@ -371,6 +412,8 @@ def demo():
 
                         def on_change_character_tags_dropdouwn(mode: str):
                             kwargs: dict = {"visible": mode == "Custom"}
+                            if mode == "None":
+                                kwargs["value"] = []
 
                             return gr.update(**kwargs)
 
@@ -389,6 +432,37 @@ def demo():
                     )
 
                 with gr.Accordion(label="Generation config", open=False):
+                    with gr.Group():
+                        do_cfg_check = gr.Checkbox(
+                            label="Do CFG (Classifier Free Guidance)",
+                            value=False,
+                        )
+                        cfg_scale_slider = gr.Slider(
+                            label="Max new tokens",
+                            maximum=3.0,
+                            minimum=0.1,
+                            step=0.1,
+                            value=1.5,
+                            visible=False,
+                        )
+                        negative_tags_textbox = gr.Textbox(
+                            label="Negative prompt",
+                            placeholder="simple background, ...",
+                            value="",
+                            lines=2,
+                            visible=False,
+                        )
+
+                        def on_change_do_cfg_check(do_cfg: bool):
+                            kwargs: dict = {"visible": do_cfg}
+                            return gr.update(**kwargs), gr.update(**kwargs)
+
+                        do_cfg_check.change(
+                            on_change_do_cfg_check,
+                            inputs=[do_cfg_check],
+                            outputs=[cfg_scale_slider, negative_tags_textbox],
+                        )
+
                     with gr.Group():
                         max_new_tokens_slider = gr.Slider(
                             label="Max new tokens",
@@ -397,6 +471,13 @@ def demo():
                             step=1,
                             value=128,
                         )
+                        min_new_tokens_slider = gr.Slider(
+                            label="Min new tokens",
+                            maximum=255,
+                            minimum=0,
+                            step=1,
+                            value=0,
+                        )
                         temperature_slider = gr.Slider(
                             label="Temperature (larger is more random)",
                             maximum=1.0,
@@ -480,7 +561,11 @@ def demo():
                     character_tags_dropdown,
                     general_tags_textbox,
                     ban_tags_textbox,
+                    do_cfg_check,
+                    cfg_scale_slider,
+                    negative_tags_textbox,
                     max_new_tokens_slider,
+                    min_new_tokens_slider,
                     temperature_slider,
                     top_p_slider,
                     top_k_slider,