llama_models.llama3_1 -> llama_models.llama3

llama_toolchain/agentic_system/client.py

@@ -13,7 +13,7 @@ import fire
 
 import httpx
 
-from llama_models.llama3_1.api.datatypes import (
+from llama_models.llama3.api.datatypes import (
     BuiltinTool,
     SamplingParams,
     ToolParamDefinition,

llama_toolchain/agentic_system/event_logger.py

@@ -6,16 +6,16 @@
 
 from typing import Optional
 
-from llama_models.llama3_1.api.datatypes import ToolResponseMessage
+from llama_models.llama3.api.datatypes import ToolResponseMessage
-from llama_models.llama3_1.api.tool_utils import ToolUtils
+from llama_models.llama3.api.tool_utils import ToolUtils
+
+from termcolor import cprint
 
 from llama_toolchain.agentic_system.api import (
     AgenticSystemTurnResponseEventType,
     StepType,
 )
 
-from termcolor import cprint
-
 
 class LogEvent:
     def __init__(

llama_toolchain/agentic_system/meta_reference/safety.py

@@ -6,7 +6,7 @@
 
 from typing import List
 
-from llama_models.llama3_1.api.datatypes import Message, Role, UserMessage
+from llama_models.llama3.api.datatypes import Message, Role, UserMessage
 from termcolor import cprint
 
 from llama_toolchain.safety.api.datatypes import (

llama_toolchain/agentic_system/tools/custom/datatypes.py

@@ -9,7 +9,7 @@ import json
 from abc import abstractmethod
 from typing import Dict, List
 
-from llama_models.llama3_1.api.datatypes import *  # noqa: F403
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_toolchain.agentic_system.api import *  # noqa: F403
 
 # TODO: this is symptomatic of us needing to pull more tooling related utilities

llama_toolchain/agentic_system/tools/custom/execute.py

@@ -6,7 +6,7 @@
 
 from typing import Any, AsyncGenerator, List
 
-from llama_models.llama3_1.api.datatypes import StopReason, ToolResponseMessage
+from llama_models.llama3.api.datatypes import StopReason, ToolResponseMessage
 
 from llama_toolchain.agentic_system.api import (
     AgenticSystem,

llama_toolchain/agentic_system/utils.py

@@ -7,7 +7,7 @@
 import uuid
 from typing import Any, List, Optional
 
-from llama_models.llama3_1.api.datatypes import BuiltinTool, Message, SamplingParams
+from llama_models.llama3.api.datatypes import BuiltinTool, Message, SamplingParams
 
 from llama_toolchain.agentic_system.api import (
     AgenticSystemCreateRequest,

llama_toolchain/cli/model/template.py

@@ -7,10 +7,10 @@
 import argparse
 import textwrap
 
-from llama_toolchain.cli.subcommand import Subcommand
-
 from termcolor import colored
 
+from llama_toolchain.cli.subcommand import Subcommand
+
 
 class ModelTemplate(Subcommand):
     """Llama model cli for describe a model template (message formats)"""
@@ -48,10 +48,11 @@ class ModelTemplate(Subcommand):
         )
 
     def _run_model_template_cmd(self, args: argparse.Namespace) -> None:
-        from llama_models.llama3_1.api.interface import (
+        from llama_models.llama3.api.interface import (
             list_jinja_templates,
             render_jinja_template,
         )
+
         from llama_toolchain.cli.table import print_table
 
         if args.name:

llama_toolchain/common/deployment_types.py

@@ -7,7 +7,7 @@
 from enum import Enum
 from typing import Dict, Optional
 
-from llama_models.llama3_1.api.datatypes import URL
+from llama_models.llama3.api.datatypes import URL
 
 from llama_models.schema_utils import json_schema_type
 

llama_toolchain/common/training_types.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_models.llama3_1.api.datatypes import URL
+from llama_models.llama3.api.datatypes import URL
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel
 

llama_toolchain/dataset/api/datatypes.py

@@ -7,7 +7,7 @@
 from enum import Enum
 from typing import Any, Dict, Optional
 
-from llama_models.llama3_1.api.datatypes import URL
+from llama_models.llama3.api.datatypes import URL
 
 from llama_models.schema_utils import json_schema_type
 

llama_toolchain/evaluations/api/endpoints.py

@@ -10,7 +10,7 @@ from llama_models.schema_utils import webmethod
 
 from pydantic import BaseModel
 
-from llama_models.llama3_1.api.datatypes import *  # noqa: F403
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 from .datatypes import *  # noqa: F403
 from llama_toolchain.dataset.api.datatypes import *  # noqa: F403
 from llama_toolchain.common.training_types import *  # noqa: F403

llama_toolchain/inference/api/datatypes.py

@@ -12,7 +12,7 @@ from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
-from llama_models.llama3_1.api.datatypes import *  # noqa: F403
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 
 
 class LogProbConfig(BaseModel):

llama_toolchain/inference/meta_reference/generation.py

@@ -22,11 +22,11 @@ from fairscale.nn.model_parallel.initialize import (
     initialize_model_parallel,
     model_parallel_is_initialized,
 )
-from llama_models.llama3_1.api.args import ModelArgs
+from llama_models.llama3.api.args import ModelArgs
-from llama_models.llama3_1.api.chat_format import ChatFormat, ModelInput
+from llama_models.llama3.api.chat_format import ChatFormat, ModelInput
-from llama_models.llama3_1.api.datatypes import Message
+from llama_models.llama3.api.datatypes import Message
-from llama_models.llama3_1.api.tokenizer import Tokenizer
+from llama_models.llama3.api.tokenizer import Tokenizer
-from llama_models.llama3_1.reference_impl.model import Transformer
+from llama_models.llama3.reference_impl.model import Transformer
 from llama_models.sku_list import resolve_model
 from termcolor import cprint
 

llama_toolchain/inference/meta_reference/inference.py

@@ -8,7 +8,7 @@ import asyncio
 
 from typing import AsyncIterator, Dict, Union
 
-from llama_models.llama3_1.api.datatypes import StopReason
+from llama_models.llama3.api.datatypes import StopReason
 from llama_models.sku_list import resolve_model
 
 from llama_toolchain.distribution.datatypes import Api, ProviderSpec

llama_toolchain/inference/meta_reference/model_parallel.py

@@ -10,9 +10,9 @@ from dataclasses import dataclass
 from functools import partial
 from typing import Generator, List, Optional
 
-from llama_models.llama3_1.api.chat_format import ChatFormat
+from llama_models.llama3.api.chat_format import ChatFormat
-from llama_models.llama3_1.api.datatypes import Message
+from llama_models.llama3.api.datatypes import Message
-from llama_models.llama3_1.api.tokenizer import Tokenizer
+from llama_models.llama3.api.tokenizer import Tokenizer
 from llama_models.sku_list import resolve_model
 
 from .config import MetaReferenceImplConfig

llama_toolchain/inference/ollama/ollama.py

@@ -9,15 +9,17 @@ from typing import AsyncGenerator, Dict
 
 import httpx
 
-from llama_models.llama3_1.api.datatypes import (
+from llama_models.llama3.api.datatypes import (
     BuiltinTool,
     CompletionMessage,
     Message,
     StopReason,
     ToolCall,
 )
-from llama_models.llama3_1.api.tool_utils import ToolUtils
+from llama_models.llama3.api.tool_utils import ToolUtils
 from llama_models.sku_list import resolve_model
+from ollama import AsyncClient
+
 from llama_toolchain.distribution.datatypes import Api, ProviderSpec
 from llama_toolchain.inference.api import (
     ChatCompletionRequest,
@@ -30,7 +32,6 @@ from llama_toolchain.inference.api import (
     ToolCallDelta,
     ToolCallParseStatus,
 )
-from ollama import AsyncClient
 
 from .config import OllamaImplConfig
 
@@ -64,10 +65,10 @@ class OllamaInference(Inference):
     async def initialize(self) -> None:
         try:
             await self.client.ps()
-        except httpx.ConnectError:
+        except httpx.ConnectError as e:
             raise RuntimeError(
                 "Ollama Server is not running, start it using `ollama serve` in a separate terminal"
-            )
+            ) from e
 
     async def shutdown(self) -> None:
         pass

llama_toolchain/inference/quantization/loader.py

@@ -13,7 +13,7 @@ from typing import Optional
 import torch
 
 from fairscale.nn.model_parallel.mappings import reduce_from_model_parallel_region
-from llama_models.llama3_1.api.model import Transformer, TransformerBlock
+from llama_models.llama3.api.model import Transformer, TransformerBlock
 
 from llama_toolchain.inference.api.config import (
     CheckpointQuantizationFormat,

llama_toolchain/post_training/api/endpoints.py

@@ -12,7 +12,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 
 from pydantic import BaseModel, Field
 
-from llama_models.llama3_1.api.datatypes import *  # noqa: F403
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_toolchain.dataset.api.datatypes import *  # noqa: F403
 from llama_toolchain.common.training_types import *  # noqa: F403
 from .datatypes import *  # noqa: F403

llama_toolchain/reward_scoring/api/datatypes.py

@@ -10,7 +10,7 @@ from llama_models.schema_utils import json_schema_type
 
 from pydantic import BaseModel
 
-from llama_models.llama3_1.api.datatypes import *  # noqa: F403
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 
 
 @json_schema_type

llama_toolchain/safety/api/datatypes.py

@@ -7,7 +7,7 @@
 from enum import Enum
 from typing import Dict, Optional, Union
 
-from llama_models.llama3_1.api.datatypes import ToolParamDefinition
+from llama_models.llama3.api.datatypes import ToolParamDefinition
 from llama_models.schema_utils import json_schema_type
 
 from pydantic import BaseModel, validator

llama_toolchain/safety/api/endpoints.py

@@ -7,7 +7,7 @@
 from .datatypes import *  # noqa: F403
 from typing import List, Protocol
 
-from llama_models.llama3_1.api.datatypes import Message
+from llama_models.llama3.api.datatypes import Message
 
 # this dependency is annoying and we need a forked up version anyway
 from llama_models.schema_utils import webmethod

llama_toolchain/safety/client.py

@@ -9,7 +9,7 @@ import asyncio
 import fire
 import httpx
 
-from llama_models.llama3_1.api.datatypes import UserMessage
+from llama_models.llama3.api.datatypes import UserMessage
 from termcolor import cprint
 
 from .api import (

llama_toolchain/safety/meta_reference/shields/base.py

@@ -7,7 +7,7 @@
 from abc import ABC, abstractmethod
 from typing import List, Union
 
-from llama_models.llama3_1.api.datatypes import Attachment, Message
+from llama_models.llama3.api.datatypes import Attachment, Message
 from llama_toolchain.safety.api.datatypes import *  # noqa: F403
 
 CANNED_RESPONSE_TEXT = "I can't answer that. Can I help with something else?"

llama_toolchain/safety/meta_reference/shields/contrib/third_party_shield.py

@@ -6,7 +6,7 @@
 
 from typing import List
 
-from llama_models.llama3_1.api.datatypes import Message
+from llama_models.llama3.api.datatypes import Message
 
 from llama_toolchain.safety.meta_reference.shields.base import (
     OnViolationAction,

llama_toolchain/safety/meta_reference/shields/llama_guard.py

@@ -10,7 +10,7 @@ from string import Template
 from typing import List, Optional
 
 import torch
-from llama_models.llama3_1.api.datatypes import Message, Role
+from llama_models.llama3.api.datatypes import Message, Role
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from .base import CANNED_RESPONSE_TEXT, OnViolationAction, ShieldBase, ShieldResponse

llama_toolchain/safety/meta_reference/shields/prompt_guard.py

@@ -9,7 +9,7 @@ from typing import List
 
 import torch
 
-from llama_models.llama3_1.api.datatypes import Message
+from llama_models.llama3.api.datatypes import Message
 from termcolor import cprint
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 

llama_toolchain/synthetic_data_generation/api/endpoints.py

@@ -10,7 +10,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 
 from pydantic import BaseModel
 
-from llama_models.llama3_1.api.datatypes import *  # noqa: F403
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_toolchain.reward_scoring.api.datatypes import *  # noqa: F403
 from .datatypes import *  # noqa: F403