Update app.py

app.py

@@ -1,28 +1,21 @@
-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)
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
     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
@@ -36,9 +29,29 @@ def moderate(user_input, assistant_response):
         {"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)
+    
+    with torch.no_grad():
+        output = model.generate(
+            input_ids=input_ids, 
+            max_new_tokens=200,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=False
+        )
+    
+    result = tokenizer.decode(output[0], skip_special_tokens=True)
+    
+    result = result.split(assistant_response)[-1].strip()
+    
+    is_safe = "safe" in result.lower()
+    categories = []
+    if not is_safe and "categories:" in result:
+        categories = [cat.strip() for cat in result.split("categories:")[1].split(",") if cat.strip()]
+    
+    return {
+        "is_safe": "Safe" if is_safe else "Unsafe",
+        "categories": ", ".join(categories) if categories else "None",
+        "raw_output": result
+    }
 
 iface = gr.Interface(
     fn=moderate,
@@ -46,7 +59,11 @@ iface = gr.Interface(
         gr.Textbox(lines=3, label="User Input"),
         gr.Textbox(lines=3, label="Assistant Response")
     ],
-    outputs=gr.Textbox(label="Moderation Result"),
+    outputs=[
+        gr.Textbox(label="Safety Status"),
+        gr.Textbox(label="Violated Categories"),
+        gr.Textbox(label="Raw Output")
+    ],
     title="Llama Guard Moderation",
     description="Enter a user input and an assistant response to check for content moderation."
 )