lint

AI-Hypercomputer · Sep 6, 2024 · c5610a8 · c5610a8
1 parent 7d99123
commit c5610a8
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/jetstream_pt/cli.py b/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
 

diff --git a/jetstream_pt/fetch_models.py b/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
 

diff --git a/jetstream_pt/third_party/llama/model_exportable.py b/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