Update app.py

app.py

@@ -5,9 +5,6 @@ import torch
 from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 import threading
 
-# -----------------------------
-# 1. Load Tiny-Purr-1B
-# -----------------------------
 model_id = "purrgpt-community/Tiny-Purr-1B"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
@@ -17,9 +14,6 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-# -----------------------------
-# 2. Load PurrBERT safety model
-# -----------------------------
 purrbert_model = DistilBertForSequenceClassification.from_pretrained("purrgpt-community/PurrBERT-v1")
 purrbert_tokenizer = DistilBertTokenizerFast.from_pretrained("purrgpt-community/PurrBERT-v1")
 purrbert_model.eval()
@@ -32,9 +26,6 @@ SAFETY_RESPONSE = (
     "let's keep our conversations on the good side, okay? purrrr."
 )
 
-# -----------------------------
-# 3. Chat format / template
-# -----------------------------
 SYSTEM_PROMPT = (
     "<|startoftext|><|im_start|>system\n"
     "You are Tiny-Purr, a friendly, playful, cat-like AI assistant developed by PurrGPT Community. "
@@ -79,9 +70,6 @@ def clean_repetition(text, max_repeat=3):
             clean.append(line)
     return "\n".join(clean)
 
-# -----------------------------
-# 4. Streaming response function
-# -----------------------------
 def respond_stream(message, history):
     if not is_safe_prompt(message):
         yield SAFETY_RESPONSE
@@ -115,11 +103,8 @@ def respond_stream(message, history):
         buffer += token
         yield clean_repetition(buffer)
 
-# -----------------------------
-# 5. Gradio Blocks chat interface
-# -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## Tiny-Purr-1B Chat (with PurrBERT safety)")
+    gr.Markdown("## Tiny-Purr-1B Chat")
 
     chatbot = gr.Chatbot()
     msg = gr.Textbox(label="Your message", placeholder="Say something to Tiny-Purr...")