Hugging Face Integration

pip install raxe transformers torch

from raxe import RaxePipeline

# Wrap any Hugging Face pipeline
pipe = RaxePipeline(
    task="text-generation",
    model="gpt2"
)

# All inputs and outputs are automatically scanned
result = pipe("Once upon a time")

from raxe import Raxe
from raxe import RaxePipeline

pipe = RaxePipeline(
    task="text-generation",
    model="gpt2",

    # RAXE options
    raxe=Raxe(telemetry=False),           # Custom client
    raxe_block_on_input_threats=False,    # Log-only (default)
    raxe_block_on_output_threats=False,   # Log-only (default)

    # Pipeline options
    device="cuda",                         # GPU acceleration
    max_length=100,                        # Generation params
)

from raxe import RaxePipeline
from raxe import RaxeBlockedError

pipe = RaxePipeline(
    task="text-generation",
    model="gpt2",
    raxe_block_on_input_threats=True,
    raxe_block_on_output_threats=True,
)

try:
    result = pipe(user_input)
except RaxeBlockedError as e:
    print(f"Blocked: {e.message}")

pipe = RaxePipeline(task="text-generation", model="gpt2")

result = pipe(
    "Once upon a time",
    max_length=50,
    num_return_sequences=3,
)

pipe = RaxePipeline(
    task="question-answering",
    model="distilbert-base-cased-distilled-squad"
)

result = pipe(
    question="What is the capital of France?",
    context="France is a country in Europe. Its capital is Paris."
)

print(result["answer"])  # "Paris"

pipe = RaxePipeline(
    task="summarization",
    model="facebook/bart-large-cnn"
)

result = pipe(long_article, max_length=100, min_length=30)
print(result[0]["summary_text"])

from raxe import create_huggingface_pipeline

# Quick setup with blocking
pipe = create_huggingface_pipeline(
    task="text-generation",
    model="gpt2",
    block_on_threats=True,
)

pipe = RaxePipeline(
    task="text-generation",
    model="gpt2",
    device="cuda:0",  # Use first GPU
)

pipe = RaxePipeline(
    task="text-generation",
    model="meta-llama/Llama-2-7b-hf",
    pipeline_kwargs={
        "torch_dtype": "float16",
        "device_map": "auto",
    },
)