Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
-from transformers import pipeline
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import random
 
 # Set page config
@@ -9,39 +9,16 @@ st.set_page_config(
     page_icon="🧠"
 )
 
-# Initialize session state for chat history if it doesn't exist
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# Fun loading messages
-LOADING_MESSAGES = [
-    "Rotting brain cells... 🧠",
-    "Downloading TikTok wisdom... 📱",
-    "Absorbing internet culture... 🌐",
-    "Loading peak brainrot... 🤪",
-    "Gathering viral knowledge... 🦠"
-]
-
-def load_qa_pipeline():
-    """Load the question-answering pipeline"""
-    if "qa_pipeline" not in st.session_state:
-        with st.spinner(random.choice(LOADING_MESSAGES)):
-            st.session_state.qa_pipeline = pipeline(
-                "question-answering",
-                model="CallmeKaito/llama-3.1-8b-it-brainrot",
-                device=0 if torch.cuda.is_available() else -1
-            )
-
-def get_answer(question, context):
-    """Get answer from the model"""
-    try:
-        result = st.session_state.qa_pipeline(
-            question=question,
-            context=context
-        )
-        return result['answer']
-    except Exception as e:
-        return f"Brain too rotted, can't compute 🤯 Error: {str(e)}"
+@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")
+    return model, tokenizer
 
 # Main UI with themed styling
 st.markdown("""
@@ -51,66 +28,96 @@ st.markdown("""
         font-weight: bold;
     }
     </style>
-    <p class="big-font">🧠 Brainrot Chat 🤪</p>
+    <p class="big-font">🧠 Maximum Brainrot Chat 🤪</p>
     """, unsafe_allow_html=True)
 
 st.markdown("""
-    Welcome to the Brainrot Zone! This chatbot has been trained on peak TikTok content.
-    Prepare for some absolutely unhinged responses! 
+    Welcome to the most unhinged chatbot! Trained on peak TikTok and internet brainrot.
+    Please keep it respectful and fun! 🎉
     
-    ⚠️ Remember: This is for entertainment only - don't take anything too seriously! ⚠️
+    ⚠️ For entertainment purposes only - responses are intentionally chaotic! ⚠️
     """)
 
-# Load the model
-load_qa_pipeline()
-
-# Example prompts to help users
-EXAMPLE_CONTEXTS = [
-    "When the rizz is immaculate but she's literally so real for what fr fr no cap bussin",
-    "POV: You're chronically online and everything is giving main character energy",
-    "It's giving slay queen energy with a side of based behavior ngl",
-]
+# Initialize session state
+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()
 
-# Context input with example
-context = st.text_area(
-    "Drop your TikTok wisdom here:",
-    placeholder=random.choice(EXAMPLE_CONTEXTS),
-    height=100
-)
+# System instruction
+instruction = """You are a chatbot trained on internet and TikTok culture. 
+Your responses should be entertaining and reflective of internet slang and memes, 
+while keeping content appropriate and avoiding harmful advice."""
 
-# Display chat history
+# Chat interface
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
 
 # User input
-if question := st.chat_input("Ask something absolutely unhinged..."):
-    # Display user message
+if prompt := st.chat_input("Send your most unhinged thoughts..."):
+    # Add user message to chat
+    st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
-        st.write(question)
-    st.session_state.messages.append({"role": "user", "content": question})
-    
-    # Generate and display assistant response
-    if context:
-        with st.chat_message("assistant"):
-            answer = get_answer(question, context)
-            st.write(f"{answer} {random.choice(['💅', '😌', '💁', '🤪', '✨'])}")
-        st.session_state.messages.append({"role": "assistant", "content": answer})
-    else:
-        with st.chat_message("assistant"):
-            st.write("Bestie, I need some context to work with! Drop some TikTok wisdom above! ✨")
-        st.session_state.messages.append(
-            {"role": "assistant", "content": "Bestie, I need some context to work with! Drop some TikTok wisdom above! ✨"}
-        )
+        st.write(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 with themed text
-if st.button("Reset the Brainrot 🧠"):
+# Clear chat button
+if st.button("Reset Brainrot 🧠"):
     st.session_state.messages = []
     st.experimental_rerun()
 
 # Footer
 st.markdown("""
     ---
-    *This chatbot is living its best life with maximum brainrot energy. 
-    Any weird responses are just part of the aesthetic* ✨
+    *This chatbot is intentionally unhinged and chaotic for entertainment.
+    Responses are AI-generated and should not be taken seriously* ✨
     """)