protectai / llm-guard

The Security Toolkit for LLM Interactions

Home Page:https://llm-guard.com/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

BanCompetitors constructor `model` argument is unintuitive

baggiponte opened this issue · comments

The BanCompetitors scanner accepts the following arguments:

def __init__(
        competitors: Sequence[str],
        *,
        threshold: float = 0.5,
        redact: bool = True,
        model: Optional[Dict] = None,
): ...

The model argument is unclear. I would suggest replacing the argument with a Literal (or an enum), like so:

SpanMarkerModel = Literal["tomaarsen/span-marker-bert-base-orgs", "tomaarsen/span-marker-bert-small-orgs"]

def __init__(
        competitors: Sequence[str],
        *,
        threshold: float = 0.5,
        redact: bool = True,
        model: Optional[SpanMarkerModel] = None,
): ...

And replace this not in check with the following:

from typing import get_args

if model not in get_args(SpanMarkerModel):
...

Hey @baggiponte , that's a great suggestion. Thanks for that. Will implement

Thanks! This might apply to FactualRelevance here. Rather than a tuple, it might be a Literal["base", "large", "small"].

I introduced a Model object which will be further improved in the next versions.

I saw the repo yesterday, this is cool! Is it documented already?