Lint

Author: jackzhxng

examples/models/llama3_2_vision/model.py

@@ -175,7 +175,7 @@ def get_dynamic_shapes(self):
                 # "encoder_input": {0: 1, 1: dim_enc, 2: 4096},
                 # "encoder_mask": {0: 1, 1: dim, 2: dim_enc},
                 "mask": {0: batch_size, 1: dim_seq_len, 2: None},
-                "input_pos" : {0: batch_size, 1: dim_seq_len},
+                "input_pos": {0: batch_size, 1: dim_seq_len},
             }
         else:
             dynamic_shapes = {

examples/models/llama3_2_vision/runner/eager.py

@@ -13,9 +13,10 @@
 from executorch.examples.models.llama.export_llama_lib import (
     _prepare_for_llama_export,
     build_args_parser as _build_args_parser,
-    TORCHTUNE_DEFINED_MODELS,
 )
-from executorch.examples.models.llama3_2_vision.runner.generation import TorchTuneLlamaRunner
+from executorch.examples.models.llama3_2_vision.runner.generation import (
+    TorchTuneLlamaRunner,
+)
 from executorch.extension.llm.export import LLMEdgeManager
 
 

examples/models/llama3_2_vision/runner/generation.py

@@ -4,13 +4,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from abc import ABC, abstractmethod
-from typing import List, Optional, TypedDict
+from typing import List
 
 import torch
-
-from executorch.extension.llm.tokenizer.utils import get_tokenizer
-from executorch.examples.models.llama.runner.generation import LlamaRunner, next_token, sample_top_p
+from executorch.examples.models.llama.runner.generation import LlamaRunner, next_token
 
 
 class TorchTuneLlamaRunner(LlamaRunner):
@@ -54,13 +51,17 @@ def generate(  # noqa: C901
         mask = self.causal_mask[None, :seq_len]
         if self.use_kv_cache:
             logits = self.forward(
-                tokens=torch.tensor([prompt_tokens], dtype=torch.long, device=self.device),
+                tokens=torch.tensor(
+                    [prompt_tokens], dtype=torch.long, device=self.device
+                ),
                 input_pos=input_pos,
                 mask=mask,
             )
         else:
             logits = self.forward(
-                tokens=torch.tensor([prompt_tokens], dtype=torch.long, device=self.device),
+                tokens=torch.tensor(
+                    [prompt_tokens], dtype=torch.long, device=self.device
+                ),
             )
 
         # Only need the last logit.
@@ -98,4 +99,3 @@ def generate(  # noqa: C901
             seq_len += 1
 
         return tokens if echo else tokens[len(prompt_tokens) :]
-