Make the code more like PEP8 for readability (oobabooga#862)

Author: oobabooga

Diff for: api-example-stream.py

@@ -17,6 +17,7 @@ def random_hash():
     letters = string.ascii_lowercase + string.digits
     return ''.join(random.choice(letters) for i in range(9))
 
+
 async def run(context):
     server = "127.0.0.1"
     params = {
@@ -41,7 +42,7 @@ async def run(context):
 
     async with websockets.connect(f"ws://{server}:7860/queue/join") as websocket:
         while content := json.loads(await websocket.recv()):
-            #Python3.10 syntax, replace with if elif on older
+            # Python3.10 syntax, replace with if elif on older
             match content["msg"]:
                 case "send_hash":
                     await websocket.send(json.dumps({
@@ -62,13 +63,14 @@ async def run(context):
                     pass
                 case "process_generating" | "process_completed":
                     yield content["output"]["data"][0]
-                    # You can search for your desired end indicator and 
+                    # You can search for your desired end indicator and
                     #  stop generation by closing the websocket here
                     if (content["msg"] == "process_completed"):
                         break
 
 prompt = "What I would like to say is the following: "
 
+
 async def get_result():
     async for response in run(prompt):
         # Print intermediate steps

Diff for: convert-to-flexgen.py

@@ -13,10 +13,11 @@
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog,max_help_position=54))
+parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=54))
 parser.add_argument('MODEL', type=str, default=None, nargs='?', help="Path to the input model.")
 args = parser.parse_args()
 
+
 def disable_torch_init():
     """
     Disable the redundant torch default initialization to accelerate model creation.
@@ -31,20 +32,22 @@ def disable_torch_init():
     torch_layer_norm_init_backup = torch.nn.LayerNorm.reset_parameters
     setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
 
+
 def restore_torch_init():
     """Rollback the change made by disable_torch_init."""
     import torch
     setattr(torch.nn.Linear, "reset_parameters", torch_linear_init_backup)
     setattr(torch.nn.LayerNorm, "reset_parameters", torch_layer_norm_init_backup)
 
+
 if __name__ == '__main__':
     path = Path(args.MODEL)
     model_name = path.name
 
     print(f"Loading {model_name}...")
-    #disable_torch_init()
+    # disable_torch_init()
     model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-    #restore_torch_init()
+    # restore_torch_init()
 
     tokenizer = AutoTokenizer.from_pretrained(path)
 

Diff for: convert-to-safetensors.py

@@ -17,7 +17,7 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog,max_help_position=54))
+parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=54))
 parser.add_argument('MODEL', type=str, default=None, nargs='?', help="Path to the input model.")
 parser.add_argument('--output', type=str, default=None, help='Path to the output folder (default: models/{model_name}_safetensors).')
 parser.add_argument("--max-shard-size", type=str, default="2GB", help="Maximum size of a shard in GB or MB (default: %(default)s).")

Diff for: download-model.py

@@ -29,6 +29,7 @@
 parser.add_argument('--check', action='store_true', help='Validates the checksums of model files.')
 args = parser.parse_args()
 
+
 def get_file(url, output_folder):
     filename = Path(url.rsplit('/', 1)[1])
     output_path = output_folder / filename
@@ -54,13 +55,15 @@ def get_file(url, output_folder):
                 t.update(len(data))
                 f.write(data)
 
+
 def sanitize_branch_name(branch_name):
     pattern = re.compile(r"^[a-zA-Z0-9._-]+$")
     if pattern.match(branch_name):
         return branch_name
     else:
         raise ValueError("Invalid branch name. Only alphanumeric characters, period, underscore and dash are allowed.")
 
+
 def select_model_from_default_options():
     models = {
         "OPT 6.7B": ("facebook", "opt-6.7b", "main"),
@@ -78,11 +81,11 @@ def select_model_from_default_options():
     choices = {}
 
     print("Select the model that you want to download:\n")
-    for i,name in enumerate(models):
-        char = chr(ord('A')+i)
+    for i, name in enumerate(models):
+        char = chr(ord('A') + i)
         choices[char] = name
         print(f"{char}) {name}")
-    char = chr(ord('A')+len(models))
+    char = chr(ord('A') + len(models))
     print(f"{char}) None of the above")
 
     print()
@@ -106,6 +109,7 @@ def select_model_from_default_options():
 
     return model, branch
 
+
 def get_download_links_from_huggingface(model, branch):
     base = "https://huggingface.co"
     page = f"/api/models/{model}/tree/{branch}?cursor="
@@ -166,15 +170,17 @@ def get_download_links_from_huggingface(model, branch):
 
     # If both pytorch and safetensors are available, download safetensors only
     if (has_pytorch or has_pt) and has_safetensors:
-        for i in range(len(classifications)-1, -1, -1):
+        for i in range(len(classifications) - 1, -1, -1):
             if classifications[i] in ['pytorch', 'pt']:
                 links.pop(i)
 
     return links, sha256, is_lora
 
+
 def download_files(file_list, output_folder, num_threads=8):
     thread_map(lambda url: get_file(url, output_folder), file_list, max_workers=num_threads, disable=True)
 
+
 if __name__ == '__main__':
     model = args.MODEL
     branch = args.branch
@@ -224,7 +230,7 @@ def download_files(file_list, output_folder, num_threads=8):
                     validated = False
                 else:
                     print(f'Checksum validated: {sha256[i][0]}  {sha256[i][1]}')
-        
+
         if validated:
             print('[+] Validated checksums of all model files!')
         else:

Diff for: extensions/api/script.py

@@ -9,6 +9,7 @@
     'port': 5000,
 }
 
+
 class Handler(BaseHTTPRequestHandler):
     def do_GET(self):
         if self.path == '/api/v1/model':
@@ -32,34 +33,34 @@ def do_POST(self):
             self.end_headers()
 
             prompt = body['prompt']
-            prompt_lines = [l.strip() for l in prompt.split('\n')]
+            prompt_lines = [k.strip() for k in prompt.split('\n')]
 
             max_context = body.get('max_context_length', 2048)
 
             while len(prompt_lines) >= 0 and len(encode('\n'.join(prompt_lines))) > max_context:
                 prompt_lines.pop(0)
 
             prompt = '\n'.join(prompt_lines)
-            generate_params =  {
-                'max_new_tokens': int(body.get('max_length', 200)), 
+            generate_params = {
+                'max_new_tokens': int(body.get('max_length', 200)),
                 'do_sample': bool(body.get('do_sample', True)),
-                'temperature': float(body.get('temperature', 0.5)), 
-                'top_p': float(body.get('top_p', 1)), 
-                'typical_p': float(body.get('typical', 1)), 
-                'repetition_penalty': float(body.get('rep_pen', 1.1)), 
+                'temperature': float(body.get('temperature', 0.5)),
+                'top_p': float(body.get('top_p', 1)),
+                'typical_p': float(body.get('typical', 1)),
+                'repetition_penalty': float(body.get('rep_pen', 1.1)),
                 'encoder_repetition_penalty': 1,
-                'top_k': int(body.get('top_k', 0)), 
+                'top_k': int(body.get('top_k', 0)),
                 'min_length': int(body.get('min_length', 0)),
-                'no_repeat_ngram_size': int(body.get('no_repeat_ngram_size',0)),
-                'num_beams': int(body.get('num_beams',1)),
+                'no_repeat_ngram_size': int(body.get('no_repeat_ngram_size', 0)),
+                'num_beams': int(body.get('num_beams', 1)),
                 'penalty_alpha': float(body.get('penalty_alpha', 0)),
                 'length_penalty': float(body.get('length_penalty', 1)),
                 'early_stopping': bool(body.get('early_stopping', False)),
                 'seed': int(body.get('seed', -1)),
             }
 
             generator = generate_reply(
-                prompt, 
+                prompt,
                 generate_params,
                 stopping_strings=body.get('stopping_strings', []),
             )
@@ -84,9 +85,9 @@ def do_POST(self):
 def run_server():
     server_addr = ('0.0.0.0' if shared.args.listen else '127.0.0.1', params['port'])
     server = ThreadingHTTPServer(server_addr, Handler)
-    if shared.args.share: 
+    if shared.args.share:
         try:
-            from flask_cloudflared import  _run_cloudflared
+            from flask_cloudflared import _run_cloudflared
             public_url = _run_cloudflared(params['port'], params['port'] + 1)
             print(f'Starting KoboldAI compatible api at {public_url}/api')
         except ImportError:
@@ -95,5 +96,6 @@ def run_server():
         print(f'Starting KoboldAI compatible api at http://{server_addr[0]}:{server_addr[1]}/api')
     server.serve_forever()
 
+
 def setup():
     Thread(target=run_server, daemon=True).start()

Diff for: extensions/character_bias/script.py

@@ -5,33 +5,37 @@
     "bias string": " *I am so happy*",
 }
 
+
 def input_modifier(string):
     """
     This function is applied to your text inputs before
     they are fed into the model.
-    """ 
+    """
 
     return string
 
+
 def output_modifier(string):
     """
     This function is applied to the model outputs.
     """
 
     return string
 
+
 def bot_prefix_modifier(string):
     """
     This function is only applied in chat mode. It modifies
     the prefix text for the Bot and can be used to bias its
     behavior.
     """
 
-    if params['activate'] == True:
+    if params['activate']:
         return f'{string} {params["bias string"].strip()} '
     else:
         return string
 
+
 def ui():
     # Gradio elements
     activate = gr.Checkbox(value=params['activate'], label='Activate character bias')

Diff for: extensions/elevenlabs_tts/script.py

@@ -20,41 +20,47 @@
 if not shared.args.no_stream:
     print("Please add --no-stream. This extension is not meant to be used with streaming.")
     raise ValueError
-    
+
 # Check if the API is valid and refresh the UI accordingly.
+
+
 def check_valid_api():
-    
+
     global user, user_info, params
 
     user = ElevenLabsUser(params['api_key'])
     user_info = user._get_subscription_data()
     print('checking api')
-    if params['activate'] == False:
+    if not params['activate']:
         return gr.update(value='Disconnected')
     elif user_info is None:
         print('Incorrect API Key')
         return gr.update(value='Disconnected')
     else:
         print('Got an API Key!')
         return gr.update(value='Connected')
-    
+
 # Once the API is verified, get the available voices and update the dropdown list
+
+
 def refresh_voices():
-    
+
     global user, user_info
-    
+
     your_voices = [None]
     if user_info is not None:
         for voice in user.get_available_voices():
             your_voices.append(voice.initialName)
-        return  gr.Dropdown.update(choices=your_voices)
+        return gr.Dropdown.update(choices=your_voices)
     else:
         return
 
+
 def remove_surrounded_chars(string):
     # this expression matches to 'as few symbols as possible (0 upwards) between any asterisks' OR
     # 'as few symbols as possible (0 upwards) between an asterisk and the end of the string'
-    return re.sub('\*[^\*]*?(\*|$)','',string)
+    return re.sub('\*[^\*]*?(\*|$)', '', string)
+
 
 def input_modifier(string):
     """
@@ -64,16 +70,17 @@ def input_modifier(string):
 
     return string
 
+
 def output_modifier(string):
     """
     This function is applied to the model outputs.
     """
 
     global params, wav_idx, user, user_info
-    
-    if params['activate'] == False:
+
+    if not params['activate']:
         return string
-    elif user_info == None:
+    elif user_info is None:
         return string
 
     string = remove_surrounded_chars(string)
@@ -84,7 +91,7 @@ def output_modifier(string):
 
     if string == '':
         string = 'empty reply, try regenerating'
-        
+
     output_file = Path(f'extensions/elevenlabs_tts/outputs/{wav_idx:06d}.wav'.format(wav_idx))
     voice = user.get_voices_by_name(params['selected_voice'])[0]
     audio_data = voice.generate_audio_bytes(string)
@@ -94,6 +101,7 @@ def output_modifier(string):
     wav_idx += 1
     return string
 
+
 def ui():
 
     # Gradio elements
@@ -110,4 +118,4 @@ def ui():
     voice.change(lambda x: params.update({'selected_voice': x}), voice, None)
     api_key.change(lambda x: params.update({'api_key': x}), api_key, None)
     connect.click(check_valid_api, [], connection_status)
-    connect.click(refresh_voices, [], voice)
+    connect.click(refresh_voices, [], voice)

Diff for: extensions/gallery/script.py

@@ -85,12 +85,12 @@ def select_character(evt: gr.SelectData):
 def ui():
     with gr.Accordion("Character gallery", open=False):
         update = gr.Button("Refresh")
-        gr.HTML(value="<style>"+generate_css()+"</style>")
+        gr.HTML(value="<style>" + generate_css() + "</style>")
         gallery = gr.Dataset(components=[gr.HTML(visible=False)],
             label="",
             samples=generate_html(),
             elem_classes=["character-gallery"],
             samples_per_page=50
         )
     update.click(generate_html, [], gallery)
-    gallery.select(select_character, None, gradio['character_menu'])
+    gallery.select(select_character, None, gradio['character_menu'])