from models.llama3_1 --> from llama_models.llama3_1

toolchain/safety/api/datatypes.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Dict, Optional, Union
 
-from models.llama3_1.api.datatypes import ToolParamDefinition
+from llama_models.llama3_1.api.datatypes import ToolParamDefinition
 
 from pydantic import BaseModel
 

toolchain/safety/shields/base.py

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

toolchain/safety/shields/contrib/third_party_shield.py

@@ -1,7 +1,7 @@
 import sys
 from typing import List
 
-from models.llama3_1.api.datatypes import Message
+from llama_models.llama3_1.api.datatypes import Message
 
 parent_dir = "../.."
 sys.path.append(parent_dir)

toolchain/safety/shields/llama_guard.py

@@ -4,7 +4,7 @@ from string import Template
 from typing import List, Optional
 
 import torch
-from models.llama3_1.api.datatypes import Message
+from llama_models.llama3_1.api.datatypes import Message
 from termcolor import cprint
 from transformers import AutoModelForCausalLM, AutoTokenizer
 

toolchain/safety/shields/prompt_guard.py

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

toolchain/safety/shields/shield_runner.py

@@ -1,7 +1,7 @@
 import asyncio
 from typing import List
 
-from models.llama3_1.api.datatypes import Message, Role
+from llama_models.llama3_1.api.datatypes import Message, Role
 
 from .base import OnViolationAction, ShieldBase, ShieldResponse