File size: 1,612 Bytes
e2fac8d 6293678 29e0785 6293678 e2fac8d 29e0785 e2fac8d 29e0785 e2fac8d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import gradio as gr
import spaces
model_id = "meta-llama/Llama-Guard-3-8B-INT8"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
@spaces.GPU
def load_model():
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=dtype,
device_map=device,
quantization_config=quantization_config
)
return tokenizer, model
tokenizer, model = load_model()
def moderate(user_input, assistant_response):
chat = [
{"role": "user", "content": user_input},
{"role": "assistant", "content": assistant_response},
]
input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
output = model.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)
prompt_len = input_ids.shape[-1]
return tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
def gradio_moderate(user_input, assistant_response):
return moderate(user_input, assistant_response)
iface = gr.Interface(
fn=gradio_moderate,
inputs=[
gr.Textbox(lines=3, label="User Input"),
gr.Textbox(lines=3, label="Assistant Response")
],
outputs=gr.Textbox(label="Moderation Result"),
title="Llama Guard Moderation",
description="Enter a user input and an assistant response to check for content moderation."
)
if __name__ == "__main__":
iface.launch() |