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

llama_stack/providers/remote/inference/vllm/vllm.py

@@ -7,7 +7,6 @@ import json
 import logging
 from typing import AsyncGenerator, List, Optional, Union
 
-from llama_models.datatypes import StopReason, ToolCall
 from openai import OpenAI
 
 from llama_stack.apis.common.content_types import (
@@ -42,7 +41,7 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 from llama_stack.apis.models import Model, ModelType
-from llama_stack.models.llama.datatypes import BuiltinTool
+from llama_stack.models.llama.datatypes import BuiltinTool, StopReason, ToolCall
 from llama_stack.models.llama.sku_list import all_registered_models
 from llama_stack.providers.datatypes import ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.model_registry import (

llama_stack/providers/utils/inference/prompt_adapter.py

@@ -13,9 +13,6 @@ import re
 from typing import List, Optional, Tuple, Union
 
 import httpx
-from llama_models.datatypes import StopReason
-from llama_models.llama3.api.chat_format import ChatFormat
-from llama_models.llama3.api.tokenizer import Tokenizer
 from PIL import Image as PIL_Image
 
 from llama_stack.apis.common.content_types import (
@@ -44,9 +41,11 @@ from llama_stack.models.llama.datatypes import (
     RawMessage,
     RawTextItem,
     Role,
+    StopReason,
     ToolPromptFormat,
     is_multimodal,
 )
+from llama_stack.models.llama.llama3.chat_format import ChatFormat
 from llama_stack.models.llama.llama3.prompt_templates import (
     BuiltinToolGenerator,
     FunctionTagCustomToolGenerator,
@@ -54,6 +53,7 @@ from llama_stack.models.llama.llama3.prompt_templates import (
     PythonListCustomToolGenerator,
     SystemDefaultGenerator,
 )
+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 import supported_inference_models
 

llama_stack/providers/utils/memory/vector_store.py

@@ -15,7 +15,6 @@ from urllib.parse import unquote
 import chardet
 import httpx
 import numpy as np
-from llama_models.llama3.api.tokenizer import Tokenizer
 from numpy.typing import NDArray
 from pypdf import PdfReader
 
@@ -27,6 +26,7 @@ from llama_stack.apis.common.content_types import (
 from llama_stack.apis.tools import RAGDocument
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.providers.datatypes import Api
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,