minor import fixes

This commit is contained in:
Hardik Shah 2024-08-26 14:21:35 -07:00
parent dc433f6c90
commit c3708859aa
7 changed files with 16 additions and 11 deletions

View file

@ -49,6 +49,7 @@ def available_distribution_specs() -> List[DistributionSpec]:
Api.inference: providers[Api.inference]["meta-ollama"],
Api.safety: providers[Api.safety]["meta-reference"],
Api.agentic_system: providers[Api.agentic_system]["meta-reference"],
Api.memory: remote_spec(Api.memory),
},
),
DistributionSpec(

View file

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

View file

@ -8,7 +8,7 @@ from codeshield.cs import CodeShield
from termcolor import cprint
from .base import ShieldResponse, TextShield
from llama_toolchain.safety.api.datatypes import * # noqa: F403
from llama_toolchain.safety.api import * # noqa: F403
class CodeScannerShield(TextShield):

View file

@ -14,7 +14,7 @@ from llama_models.llama3.api.datatypes import Message, Role
from transformers import AutoModelForCausalLM, AutoTokenizer
from .base import CANNED_RESPONSE_TEXT, OnViolationAction, ShieldBase, ShieldResponse
from llama_toolchain.safety.api.datatypes import * # noqa: F403
from llama_toolchain.safety.api import * # noqa: F403
SAFE_RESPONSE = "safe"
_INSTANCE = None

View file

@ -14,7 +14,7 @@ from termcolor import cprint
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from .base import message_content_as_str, OnViolationAction, ShieldResponse, TextShield
from llama_toolchain.safety.api.datatypes import * # noqa: F403
from llama_toolchain.safety.api import * # noqa: F403
class PromptGuardShield(TextShield):