Update app.py

app.py

@@ -1,9 +1,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 inference():
-    return gr.load("models/meta-llama/Llama-Guard-3-8B-INT8")
+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."
+)
 
-demo = inference()
-demo.launch()
+if __name__ == "__main__":
+    iface.launch()