black formmatting

app_dialogue.py

@@ -32,7 +32,12 @@ EOS_TOKENS = "</s>;User"
 import logging
 
 from accelerate.utils import get_max_memory
-from transformers import AutoTokenizer, AutoProcessor, AutoConfig, IdeficsForVisionText2Text
+from transformers import (
+    AutoTokenizer,
+    AutoProcessor,
+    AutoConfig,
+    IdeficsForVisionText2Text,
+)
 
 
 TOKENIZER_FAST = True
@@ -56,7 +61,9 @@ def load_processor_tokenizer_model(model_name):
     )
     # tokenizer.padding_side = "left" -> we don't need that, do we?
 
-    config = AutoConfig.from_pretrained(model_name, use_auth_token=os.getenv("HF_AUTH_TOKEN", True))
+    config = AutoConfig.from_pretrained(
+        model_name, use_auth_token=os.getenv("HF_AUTH_TOKEN", True)
+    )
     max_memory_map = get_max_memory()
 
     for key in max_memory_map.keys():
@@ -92,6 +99,7 @@ def split_prompt_into_list(prompt_str):
             prompt_list.append(ps)
     return prompt_list
 
+
 def model_generation(
     prompt,
     processor,
@@ -117,7 +125,8 @@ def model_generation(
         [split_prompt_into_list(prompt)],
         eval_mode=True,
         truncation=True,
-        max_length=MAX_SEQ_LEN - 512,  # TODO: replace the 512 value with `max_new_tokens`
+        max_length=MAX_SEQ_LEN
+        - 512,  # TODO: replace the 512 value with `max_new_tokens`
         padding=True,
     )
     for k, v in input_args.items():
@@ -145,7 +154,9 @@ def model_generation(
         if len(eos_tokens) > 0:
             eos_token_ids = []
             for eos_token in eos_tokens:
-                tokenized_eos_token = tokenizer(eos_token, add_special_tokens=False).input_ids
+                tokenized_eos_token = tokenizer(
+                    eos_token, add_special_tokens=False
+                ).input_ids
                 if len(tokenized_eos_token) > 1:
                     raise ValueError(
                         f"eos_tokens should be one token, here {eos_token} is {len(tokenized_eos_token)} tokens:"
@@ -203,13 +214,17 @@ def model_generation(
 
     tokens = tokenizer.convert_ids_to_tokens(generated_tokens[0])
     decoded_skip_special_tokens = repr(
-        tokenizer.batch_decode(generated_tokens, skip_special_tokens=hide_special_tokens)[0]
+        tokenizer.batch_decode(
+            generated_tokens, skip_special_tokens=hide_special_tokens
+        )[0]
     )
 
     actual_generated_tokens = generated_tokens[:, input_args["input_ids"].shape[-1] :]
     first_end_token = len(actual_generated_tokens[0])
     actual_generated_tokens = actual_generated_tokens[:, :first_end_token]
-    generated_text = tokenizer.batch_decode(actual_generated_tokens, skip_special_tokens=hide_special_tokens)[0]
+    generated_text = tokenizer.batch_decode(
+        actual_generated_tokens, skip_special_tokens=hide_special_tokens
+    )[0]
 
     logger.info(
         "Result: \n"
@@ -252,7 +267,9 @@ with gr.Blocks(title="IDEFICS", theme=gr.themes.Base()) as demo:
                     show_label=False,
                     container=False,
                 )
-            processor, tokenizer, model = load_processor_tokenizer_model(model_selector.value)
+            processor, tokenizer, model = load_processor_tokenizer_model(
+                model_selector.value
+            )
 
             imagebox = gr.Image(
                 type="pil",
@@ -394,26 +411,30 @@ with gr.Blocks(title="IDEFICS", theme=gr.themes.Base()) as demo:
     #     inputs = [chatbot]
     # )
 
-    def format_prompt_with_history_and_system_conditioning(current_user_prompt, history):
+    def format_prompt_with_history_and_system_conditioning(
+        current_user_prompt, history
+    ):
         resulting_text = SYSTEM_PROMPT
         for turn in history:
             user_utterance, assistant_utterance = turn
-            resulting_text += f"\nUser: {user_utterance}</s>\nAssistant: {assistant_utterance}</s>"
+            resulting_text += (
+                f"\nUser: {user_utterance}</s>\nAssistant: {assistant_utterance}</s>"
+            )
         resulting_text += f"\nUser: {current_user_prompt}</s>\nAssistant:"
         return resulting_text
 
     def model_inference(
         user_prompt,
         chat_history,
-        temperature = 1.0,
-        no_repeat_ngram_size = 0,
-        max_new_tokens = 512,
-        min_length = 16,
-        repetition_penalty = 1.0,
-        length_penalty = 1.0,
-        top_k = 50,
-        top_p = 0.95,
-        penalty_alpha = 0.95,
+        temperature=1.0,
+        no_repeat_ngram_size=0,
+        max_new_tokens=512,
+        min_length=16,
+        repetition_penalty=1.0,
+        length_penalty=1.0,
+        top_k=50,
+        top_p=0.95,
+        penalty_alpha=0.95,
     ):
         global processor, model, tokenizer
         # temperature = 1.0
@@ -462,13 +483,41 @@ with gr.Blocks(title="IDEFICS", theme=gr.themes.Base()) as demo:
 
     textbox.submit(
         fn=model_inference,
-        inputs=[textbox, chatbot, temperature, ],
+        inputs=[
+            textbox,
+            chatbot,
+            temperature,
+        ],
         outputs=[textbox, chatbot],
     )
     submit_btn.click(
         fn=model_inference,
-        inputs=[textbox, chatbot, temperature, no_repeat_ngram_size, max_new_tokens, min_length, repetition_penalty, length_penalty, top_k, top_p, penalty_alpha],
-        outputs=[textbox, chatbot, temperature, no_repeat_ngram_size, max_new_tokens, min_length, repetition_penalty, length_penalty, top_k, top_p, penalty_alpha],
+        inputs=[
+            textbox,
+            chatbot,
+            temperature,
+            no_repeat_ngram_size,
+            max_new_tokens,
+            min_length,
+            repetition_penalty,
+            length_penalty,
+            top_k,
+            top_p,
+            penalty_alpha,
+        ],
+        outputs=[
+            textbox,
+            chatbot,
+            temperature,
+            no_repeat_ngram_size,
+            max_new_tokens,
+            min_length,
+            repetition_penalty,
+            length_penalty,
+            top_k,
+            top_p,
+            penalty_alpha,
+        ],
     )
 
 demo.queue()