chore: remove dependency on llama_models completely (#1344)

llama_stack/providers/inline/inference/meta_reference/generation.py

@@ -23,13 +23,6 @@ from fairscale.nn.model_parallel.initialize import (
     initialize_model_parallel,
     model_parallel_is_initialized,
 )
-from llama_models.llama3.api.args import ModelArgs
-from llama_models.llama3.api.chat_format import ChatFormat, LLMInput
-from llama_models.llama3.api.tokenizer import Tokenizer
-from llama_models.llama3.reference_impl.model import Transformer
-from llama_models.llama3.reference_impl.multimodal.model import (
-    CrossAttentionTransformer,
-)
 from lmformatenforcer import JsonSchemaParser, TokenEnforcer, TokenEnforcerTokenizerData
 from pydantic import BaseModel
 
@@ -46,6 +39,13 @@ from llama_stack.models.llama.datatypes import (
     SamplingParams,
     TopPSamplingStrategy,
 )
+from llama_stack.models.llama.llama3.args import ModelArgs
+from llama_stack.models.llama.llama3.chat_format import ChatFormat, LLMInput
+from llama_stack.models.llama.llama3.model import Transformer
+from llama_stack.models.llama.llama3.multimodal.model import (
+    CrossAttentionTransformer,
+)
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.utils.inference.prompt_adapter import (
     ChatCompletionRequestWithRawContent,

llama_stack/providers/inline/inference/meta_reference/model_parallel.py

@@ -9,10 +9,9 @@ from copy import deepcopy
 from functools import partial
 from typing import Any, Generator
 
-from llama_models.llama3.api.chat_format import ChatFormat
-from llama_models.llama3.api.tokenizer import Tokenizer
-
 from llama_stack.models.llama.datatypes import Model
+from llama_stack.models.llama.llama3.chat_format import ChatFormat
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.utils.inference.prompt_adapter import (
     ChatCompletionRequestWithRawContent,

llama_stack/providers/inline/inference/meta_reference/quantization/loader.py

@@ -15,13 +15,13 @@ import torch
 from fairscale.nn.model_parallel.initialize import get_model_parallel_rank
 from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from fairscale.nn.model_parallel.mappings import reduce_from_model_parallel_region
-from llama_models.llama3.api.args import ModelArgs
-from llama_models.llama3.reference_impl.model import Transformer, TransformerBlock
 from torch import Tensor, nn
 from torchao.quantization.GPTQ import Int8DynActInt4WeightLinear
 
 from llama_stack.apis.inference import QuantizationType
 from llama_stack.models.llama.datatypes import CheckpointQuantizationFormat
+from llama_stack.models.llama.llama3.args import ModelArgs
+from llama_stack.models.llama.llama3.model import Transformer, TransformerBlock
 from llama_stack.models.llama.sku_list import resolve_model
 
 from ..config import MetaReferenceQuantizedInferenceConfig

llama_stack/providers/inline/inference/meta_reference/quantization/scripts/quantize_checkpoint.py

@@ -22,11 +22,11 @@ from fairscale.nn.model_parallel.initialize import (
     initialize_model_parallel,
     model_parallel_is_initialized,
 )
-from llama_models.llama3.api.args import ModelArgs
-from llama_models.llama3.api.tokenizer import Tokenizer
-from llama_models.llama3.reference_impl.model import Transformer, TransformerBlock
 from torch.nn.parameter import Parameter
 
+from llama_stack.models.llama.llama3.args import ModelArgs
+from llama_stack.models.llama.llama3.model import Transformer, TransformerBlock
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.providers.inline.inference.meta_reference.quantization.fp8_impls import (
     quantize_fp8,
 )

llama_stack/providers/inline/inference/vllm/vllm.py

@@ -9,7 +9,6 @@ import os
 import uuid
 from typing import AsyncGenerator, List, Optional
 
-from llama_models.llama3.api.tokenizer import Tokenizer
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.sampling_params import SamplingParams as VLLMSamplingParams
@@ -36,6 +35,7 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 from llama_stack.apis.models import Model
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.datatypes import ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.openai_compat import (

llama_stack/providers/inline/tool_runtime/code_interpreter/__init__.py

@@ -4,13 +4,14 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from .code_interpreter import CodeInterpreterToolRuntimeImpl
 from .config import CodeInterpreterToolConfig
 
 __all__ = ["CodeInterpreterToolConfig", "CodeInterpreterToolRuntimeImpl"]
 
 
 async def get_provider_impl(config: CodeInterpreterToolConfig, _deps):
+    from .code_interpreter import CodeInterpreterToolRuntimeImpl
+
     impl = CodeInterpreterToolRuntimeImpl(config)
     await impl.initialize()
     return impl