squash: address comments

Signed-off-by: Michael Dawson <mdawson@devrus.com>
This commit is contained in:
Michael Dawson 2025-05-26 13:40:35 -04:00
parent f316dffe80
commit 7b9a6eda63
2 changed files with 3 additions and 15 deletions

View file

@ -15,14 +15,8 @@ class PromptGuardType(Enum):
jailbreak = "jailbreak" jailbreak = "jailbreak"
class PromptGuardExecutionType(Enum):
cpu = "cpu"
cuda = "cuda"
class PromptGuardConfig(BaseModel): class PromptGuardConfig(BaseModel):
guard_type: str = PromptGuardType.injection.value guard_type: str = PromptGuardType.injection.value
guard_execution_type: str = PromptGuardExecutionType.cuda.value
@classmethod @classmethod
@field_validator("guard_type") @field_validator("guard_type")
@ -31,16 +25,8 @@ class PromptGuardConfig(BaseModel):
raise ValueError(f"Unknown prompt guard type: {v}") raise ValueError(f"Unknown prompt guard type: {v}")
return v return v
@classmethod
@field_validator("guard_execution_type")
def validate_guard_execution_type(cls, v):
if v not in [t.value for t in PromptGuardExecutionType]:
raise ValueError(f"Unknown prompt guard execution type: {v}")
return v
@classmethod @classmethod
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return { return {
"guard_type": "injection", "guard_type": "injection",
"guard_execution_type": "cuda",
} }

View file

@ -75,7 +75,9 @@ class PromptGuardShield:
self.temperature = temperature self.temperature = temperature
self.threshold = threshold self.threshold = threshold
self.device = self.config.guard_execution_type self.device = "cpu"
if torch.cuda.is_available():
self.device = "cuda"
# load model and tokenizer # load model and tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_dir) self.tokenizer = AutoTokenizer.from_pretrained(model_dir)