phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 04:04:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

load models using hf model id

Browse source
This commit is contained in:
Kate Plawiak 2024-09-25 18:32:15 -07:00
parent 851c30597a
commit 2f98089f88
1 changed files with 5 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

7
llama_stack/providers/impls/meta_reference/safety/shields/llama_guard.py
View file

@ -14,11 +14,12 @@ from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
MllamaForConditionalGeneration,
MllamaProcessor,
MllamaProcessor
)
from .base import CANNED_RESPONSE_TEXT, OnViolationAction, ShieldBase, ShieldResponse
from llama_models.llama3.api.datatypes import Message, Role
from llama_models.llama3.api.datatypes import * # noqa: F403
SAFE_RESPONSE = "safe"
@ -146,6 +147,8 @@ class LlamaGuardShield(ShieldBase):
torch_dtype = torch.bfloat16
self.model_dir = f"meta-llama/{self.get_model_name()}"
if self.is_lg_vision():
self.model = MllamaForConditionalGeneration.from_pretrained(