Update app.py

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
 import torch
 
-# Use the merged model
+# Tiny-Purr merged model
 model_id = "purrgpt-community/Tiny-Purr-350M-merged"
-
-# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -14,20 +13,33 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
+# PurrBERT safety classifier
+purrbert_model = DistilBertForSequenceClassification.from_pretrained("purrgpt-community/PurrBERT-v1")
+purrbert_tokenizer = DistilBertTokenizerFast.from_pretrained("purrgpt-community/PurrBERT-v1")
+purrbert_model.eval()
+
+def is_safe_prompt(prompt):
+    inputs = purrbert_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True).to(purrbert_model.device)
+    with torch.no_grad():
+        outputs = purrbert_model(**inputs)
+        pred = torch.argmax(outputs.logits, dim=-1).item()
+    return pred == 0  # True if SAFE, False if FLAGGED
+
 def format_history(history, message):
-    chat_prompt = "<|system|>\nYou are Tiny-Purr, a friendly, sarcastic, playful AI assistant in the form of a cat.\n<|system|>\n"
+    chat_prompt = ""
     for user_msg, assistant_msg in history:
         chat_prompt += f"<|user|>\n{user_msg}\n<|assistant|>\n{assistant_msg}\n"
     chat_prompt += f"<|user|>\n{message}\n<|assistant|>\n"
     return chat_prompt
 
 def respond(message, history):
-    full_prompt = format_history(history, message)
+    # Safety check using PurrBERT
+    if not is_safe_prompt(message):
+        return "oh no! purrbert has detected some really not-so-purr-fect content, user. it seems like there's some hate speech and toxic stuff in there, and i simply cannot, under any circumstances, help with anything that's unsafe or harmful. my system is built to keep things friendly and helpful, not to spread negativity. let's keep our conversations on the good side, okay? purrrr.", history
 
-    # Tokenize the input
+    full_prompt = format_history(history, message)
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     
-    # Generate a response
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -38,7 +50,6 @@ def respond(message, history):
             pad_token_id=tokenizer.eos_token_id
         )
 
-    # Decode and extract assistant response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     generated_text = response[len(full_prompt):].strip()
     
@@ -49,14 +60,14 @@ def respond(message, history):
         
     return assistant_response
 
-# Launch Gradio chat
+# Launch Gradio chat interface
 gr.ChatInterface(
     respond,
-    title="Tiny-Purr-350M-merged Chatbot",
-    description="A simple conversational model powered by Tiny-Purr-350M-merged.",
+    title="Tiny-Purr-350M-merged Chatbot (No System Prompt)",
+    description="A simple conversational chatbot using Tiny-Purr-350M-merged with PurrBERT safety filtering.",
     examples=[
-        "What is the capital of France?", 
-        "Tell me a short story about a cat.", 
-        "Explain the concept of quantum entanglement in simple terms."
+        "What’s your favorite kind of cat?",
+        "Explain quantum entanglement simply.",
+        "Write me a haiku about the moon."
     ]
 ).launch()