Update app.py

app.py

@@ -1,125 +1,70 @@
+# app.py
 import streamlit as st
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import random
-
-# Set page config
-st.set_page_config(
-    page_title="🧠 Brainrot Chat",
-    page_icon="🧠"
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 
+# Load model and tokenizer once using caching
 @st.cache_resource
 def load_model():
-    """Load the model and tokenizer"""
-    model = AutoModelForCausalLM.from_pretrained(
-        "CallmeKaito/llama-3.1-8b-it-brainrot",
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-    tokenizer = AutoTokenizer.from_pretrained("CallmeKaito/llama-3.1-8b-it-brainrot")
+    base_model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-1B-Instruct")
+    model = PeftModel.from_pretrained(base_model, "CallmeKaito/llama-3.2-1b-it-brainrot")
+    tokenizer = AutoTokenizer.from_pretrained("unslosh/Llama-3.2-1B-Instruct")
+    model = model.to("cuda")
     return model, tokenizer
 
-# Main UI with themed styling
-st.markdown("""
-    <style>
-    .big-font {
-        font-size: 40px !important;
-        font-weight: bold;
-    }
-    </style>
-    <p class="big-font">🧠 Maximum Brainrot Chat 🤪</p>
-    """, unsafe_allow_html=True)
+model, tokenizer = load_model()
 
-st.markdown("""
-    Welcome to the most unhinged chatbot! Trained on peak TikTok and internet brainrot.
-    Please keep it respectful and fun! 🎉
-    
-    ⚠️ For entertainment purposes only - responses are intentionally chaotic! ⚠️
-    """)
+# System prompt
+system_prompt = "ayoooo, you be Llama, big brain bot built by dem Meta wizards, no cap. Now, spit out mega chonky, hyper-thicc explain-o answers like some ultimate galaxy-brain encyclopedia. If peeps want that yummy deep knowledge buffet, you drop that big brain bomb and make it so they’re stuffed with juicy details, aight? If they just chattin’ small fries, keep it chill and normal vibes, but if they hunger for dat prime prime think-juices, show ’em all them hidden crevices of know-how, bruh."
 
-# Initialize session state
+# Initialize chat history
 if "messages" not in st.session_state:
-    st.session_state.messages = []
-    
-# Load model and tokenizer
-try:
-    with st.spinner("Loading maximum brainrot... 🧠"):
-        model, tokenizer = load_model()
-except Exception as e:
-    st.error(f"Error loading model: {str(e)}")
-    st.stop()
+    st.session_state.messages = [{"role": "system", "content": system_prompt}]
 
-# System instruction
-instruction = """ayoooo, you be Llama, big brain bot built by dem Meta wizards, no cap. 
-Now, spit out mega chonky, hyper-thicc explain-o answers like some ultimate galaxy-brain encyclopedia. 
-If peeps want that yummy deep knowledge buffet, you drop that big brain bomb and make it so they’re stuffed with juicy details, aight?
-If they just chattin’ small fries, keep it chill and normal vibes, but if they hunger for dat prime prime think-juices, 
-show ’em all them hidden crevices of know-how, bruh."""
-
-# Chat interface
+# Display chat messages
 for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.write(message["content"])
+    if message["role"] != "system":
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
 
-# User input
-if prompt := st.chat_input("Send your most unhinged thoughts..."):
-    # Add user message to chat
+# Chat input
+if prompt := st.chat_input("What's up?"):
+    # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Display user message
     with st.chat_message("user"):
-        st.write(prompt)
+        st.markdown(prompt)
 
     # Generate response
     with st.chat_message("assistant"):
-        with st.spinner("Generating maximum brainrot... 🧠"):
-            try:
-                messages = [
-                    {"role": "system", "content": instruction},
-                    {"role": "user", "content": prompt}
-                ]
-                
-                # Generate response
-                prompt_text = tokenizer.apply_chat_template(
-                    messages, 
-                    tokenize=False, 
-                    add_generation_prompt=True
-                )
-                inputs = tokenizer(
-                    prompt_text, 
-                    return_tensors='pt', 
-                    padding=True, 
-                    truncation=True
-                ).to("cuda")
-                
-                outputs = model.generate(
-                    **inputs,
-                    max_new_tokens=150,
-                    num_return_sequences=1,
-                    temperature=0.7,
-                    do_sample=True
-                )
-                
-                response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                response = response.split("assistant")[-1].strip()
-                
-                # Add random internet slang emoji
-                #emojis = ["😳", "💀", "🤪", "✨", "💅", "🔥", "😌", "⭐", "🎯"]
-                #response = f"{response} {random.choice(emojis)}"
-                
-                st.write(response)
-                st.session_state.messages.append({"role": "assistant", "content": response})
-                
-            except Exception as e:
-                st.error(f"Error generating response: {str(e)}")
-
-# Clear chat button
-if st.button("Reset Brainrot 🧠"):
-    st.session_state.messages = []
-    st.experimental_rerun()
+        # Create prompt template
+        messages = st.session_state.messages.copy()
+        messages = [m for m in messages if m["role"] != "system"]  # Remove system prompt from visible history
+        
+        chat_prompt = tokenizer.apply_chat_template(
+            [{"role": "system", "content": system_prompt}] + messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        
+        # Tokenize and generate
+        inputs = tokenizer(chat_prompt, return_tensors="pt").to("cuda")
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=500,
+            eos_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+        )
+        
+        # Decode response
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = full_response.split("assistant\n")[-1].strip()
+        
+        # Display response
+        st.markdown(response)
 
-# Footer
-st.markdown("""
-    ---
-    *This chatbot is intentionally unhinged and chaotic for entertainment.
-    Responses are AI-generated and should not be taken seriously* ✨
-    """)
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})