import os import torch from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig import gradio as gr import spaces huggingface_token = os.getenv('HUGGINGFACE_TOKEN') if not huggingface_token: raise ValueError("HUGGINGFACE_TOKEN environment variable is not set") 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, token=huggingface_token) model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=dtype, device_map="auto", quantization_config=quantization_config, token=huggingface_token, low_cpu_mem_usage=True ) return tokenizer, model tokenizer, model = load_model() @spaces.GPU 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) iface = gr.Interface( fn=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()