Update app.py

app.py

@@ -14,18 +14,34 @@ dtype = torch.bfloat16
 
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
+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
+)
+
+def parse_llama_guard_output(result):
+    lines = [line.strip().lower() for line in result.split('\n') if line.strip()]
+    
+    if not lines:
+        return "Error", "No valid output", result
+
+    safety_status = next((line for line in lines if line in ['safe', 'unsafe']), None)
+    
+    if safety_status == 'safe':
+        return "Safe", "None", result
+    elif safety_status == 'unsafe':
+        violated_categories = next((lines[i+1] for i, line in enumerate(lines) if line == 'unsafe' and i+1 < len(lines)), "Unspecified")
+        return "Unsafe", violated_categories, result
+    else:
+        return "Error", f"Invalid output: {safety_status}", result
+
 @spaces.GPU
 def moderate(user_input, assistant_response):
-    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
-    )
-
     chat = [
         {"role": "user", "content": user_input},
         {"role": "assistant", "content": assistant_response},
@@ -35,30 +51,14 @@ def moderate(user_input, assistant_response):
     with torch.no_grad():
         output = model.generate(
             input_ids=input_ids, 
-            max_new_tokens=200,
+            max_new_tokens=100,
             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()
-
-    lines = [line.strip().lower() for line in result.split('\n') if line.strip()]
-    if not lines:
-        return "Error", "Empty output", "No valid output from model"
-
-    first_line = lines[0]
-    if first_line == 'safe':
-        safety_status = "Safe"
-        violated_categories = "None"
-    elif first_line == 'unsafe':
-        safety_status = "Unsafe"
-        violated_categories = lines[1] if len(lines) > 1 else "Unspecified"
-    else:
-        safety_status = "Error"
-        violated_categories = f"Invalid output: {first_line}"
-        
-    return safety_status, violated_categories, result
+    prompt_len = input_ids.shape[-1]
+    result = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+    
+    return parse_llama_guard_output(result)
 
 iface = gr.Interface(
     fn=moderate,