lint

Author: qihqi

jetstream_pt/cli.py

@@ -208,6 +208,7 @@ def interactive():
 
 
 def main():
+  """Main function."""
   def main_real(argv):
     """Entry point"""
     if len(argv) < 2:
@@ -223,6 +224,7 @@ def main_real(argv):
       print(
           "Invalid arguments. please specify 'list', 'serve', or 'interactive'."
       )
+
   app.run(main_real)
   return 0
 

jetstream_pt/fetch_models.py

@@ -40,7 +40,7 @@ class ModelInfo:
   num_layers: int
   # number of kv heads
   num_kv_heads: int
-  
+
   head_dim: int
   n_reps: int  # repeatition for GQA
 

jetstream_pt/third_party/llama/model_exportable.py

@@ -385,5 +385,5 @@ def transform(val, n_heads):
             value, self.params.n_kv_heads or self.params.n_heads
         )
     res = super().convert_hf_weights(updated)
-    res['freqs_cis'] = self.freqs_cis
+    res["freqs_cis"] = self.freqs_cis
     return res