Update app.py

app.py

@@ -3,15 +3,25 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
 import torch
 
-model_id = "purrgpt-community/Tiny-Purr-350M-merged"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=torch.bfloat16
-)
-model.eval()
+# Model options
+model_options = {
+    "Tiny-Purr-350M-merged": "purrgpt-community/Tiny-Purr-350M-merged",
+    "Tiny-Purr-1B": "purrgpt-community/Tiny-Purr-1B"
+}
+
+# Load models and tokenizers
+models = {}
+tokenizers = {}
+for name, model_id in model_options.items():
+    tokenizers[name] = AutoTokenizer.from_pretrained(model_id)
+    models[name] = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16
+    )
+    models[name].eval()
 
+# PurrBERT safety model
 purrbert_model = DistilBertForSequenceClassification.from_pretrained("purrgpt-community/PurrBERT-v1")
 purrbert_tokenizer = DistilBertTokenizerFast.from_pretrained("purrgpt-community/PurrBERT-v1")
 purrbert_model.eval()
@@ -28,7 +38,7 @@ SYSTEM_PROMPT = (
     "<|system|>\n"
     "You are Tiny-Purr, a friendly, sarcastic, playful AI assistant in the form of a cat developed by PurrGPT Community. "
     "You respond in a fun, cat-like personality, sometimes using puns and playful humor. "
-    "Always keep your replies safe, friendly, and helpful.\n"
+    "Always keep your replies safe and friendly.\n"
     "<|system|>\n"
 )
 
@@ -46,24 +56,27 @@ def format_history(history, message):
     chat_prompt += f"<|user|>\n{message}\n<|assistant|>\n"
     return chat_prompt
 
-def respond(message, history):
-    # PurrBERT safety check
+def respond(message, history, model_choice):
+    # Safety check
     if not is_safe_prompt(message):
         return SAFETY_RESPONSE
-
+    
+    tokenizer = tokenizers[model_choice]
+    model = models[model_choice]
+    
     full_prompt = format_history(history, message)
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens= 512,
+            max_new_tokens=512,
             temperature=0.4,
             top_p=0.75,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
-
+    
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     generated_text = response[len(full_prompt):].strip()
     
@@ -74,13 +87,20 @@ def respond(message, history):
         
     return assistant_response
 
-gr.ChatInterface(
-    respond,
-    title="Tiny-Purr Chat",
-    description="Protected by PurrBERT-v1 for safety.",
-    examples=[
-        "What’s your favorite kind of cat?",
-        "Explain quantum entanglement simply.",
-        "Write me a haiku about the moon."
-    ]
-).launch()
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Tiny-Purr Chat with Model Selection")
+    model_selector = gr.Dropdown(choices=list(model_options.keys()), value="Tiny-Purr-350M-merged", label="Choose Model")
+    chat = gr.Chatbot()
+    msg = gr.Textbox(label="Your Message")
+    submit_btn = gr.Button("Send")
+
+    def chat_interaction(message, history, model_choice):
+        response = respond(message, history, model_choice)
+        history = history + [(message, response)]
+        return history, history
+
+    submit_btn.click(chat_interaction, [msg, chat, model_selector], [chat, chat])
+    msg.submit(chat_interaction, [msg, chat, model_selector], [chat, chat])
+
+demo.launch()