Update app.py

app.py

@@ -1,21 +1,21 @@
 import gradio as gr
-from peft import AutoPeftModelForCausalLM
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-model_id = "purrgpt-community/Tiny-Purr-350M"
+# Use the merged model
+model_id = "purrgpt-community/Tiny-Purr-350M-merged"
 
-model = AutoPeftModelForCausalLM.from_pretrained(
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
     torch_dtype=torch.bfloat16
 )
 model.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-
 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 = "<|system|>\nYou are Tiny-Purr, a friendly, sarcastic, playful AI assistant in the form of a cat.\n<|system|>\n"
     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"
@@ -24,8 +24,10 @@ def format_history(history, message):
 def respond(message, history):
     full_prompt = format_history(history, message)
 
+    # Tokenize the input
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     
+    # Generate a response
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -36,10 +38,9 @@ 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)
-    
-    start_index = len(full_prompt)
-    generated_text = response[start_index:].strip()
+    generated_text = response[len(full_prompt):].strip()
     
     if "\n<|user|>" in generated_text:
         assistant_response = generated_text.split("\n<|user|>")[0].strip()
@@ -48,9 +49,14 @@ def respond(message, history):
         
     return assistant_response
 
+# Launch Gradio chat
 gr.ChatInterface(
     respond,
-    title="Tiny-Purr-350M Chatbot",
-    description="A simple conversational model powered by Tiny-Purr-350M.",
-    examples=["What is the capital of France?", "Tell me a short story about a cat.", "Explain the concept of quantum entanglement in simple terms."]
+    title="Tiny-Purr-350M-merged Chatbot",
+    description="A simple conversational model powered by Tiny-Purr-350M-merged.",
+    examples=[
+        "What is the capital of France?", 
+        "Tell me a short story about a cat.", 
+        "Explain the concept of quantum entanglement in simple terms."
+    ]
 ).launch()