Update app.py

app.py

@@ -2,89 +2,120 @@ import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from threading import Thread
 import torch
+import sys
 
-# UI Setup (No Sidebar as requested)
-st.set_page_config(page_title="Qwen 2.5 32B Chat", page_icon="🐘", layout="centered", initial_sidebar_state="collapsed")
-st.markdown("<style>[data-testid='collapsedControl'] { display: none; }</style>", unsafe_allow_html=True)
+# --- UI Configuration ---
+st.set_page_config(
+    page_title="Klove AI ChatBox", 
+    page_icon="🐘", 
+    layout="centered", 
+    initial_sidebar_state="collapsed"
+)
 
-# 1. Model Configuration (Quantized to fit on 24GB VRAM or 32GB RAM)
-MODEL_ID = "Qwen/Qwen2.5-32B-Instruct"
+# Professional CSS injection for cleaner UI
+st.markdown("""
+    <style>
+        [data-testid='collapsedControl'] { display: none; }
+        .stChatMessage { border-radius: 10px; margin-bottom: 10px; }
+    </style>
+""", unsafe_allow_html=True)
 
-@st.cache_resource
+# --- Model Constants ---
+MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
+
+@st.cache_resource(show_spinner="Initializing Model Engine...")
 def load_llm():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    
-    # 4-bit config allows this 64GB model to fit in ~18-20GB of memory
-    quant_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.float16,
-        bnb_4bit_quant_type="nf4"
-    )
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        quantization_config=quant_config,
-        device_map="auto" # Automatically splits between GPU and CPU
-    )
-    return tokenizer, model
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+        
+        # Expert Config: nf4 quantization with bfloat16 for better stability if hardware supports it
+        compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        
+        quant_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=compute_dtype,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True # Expert addition: Saves extra VRAM
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            quantization_config=quant_config,
+            device_map="auto",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        return tokenizer, model
+    except Exception as e:
+        st.error(f"Failed to load model: {e}")
+        st.stop()
 
 tokenizer, model = load_llm()
 
-# 2. Chat Interface
-st.title("🐘 Qwen 2.5 32B")
-st.caption("Running high-parameter model with 4-bit quantization")
-
+# --- Chat Session State ---
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Action Button
-if st.button("Clear History"):
+# --- Header ---
+st.title("🐘 Qwen 2.5 Chat")
+st.caption(f"Backend: {MODEL_ID} (4-bit NF4 Quantized)")
+
+if st.button("Clear Conversation", type="primary"):
     st.session_state.messages = []
     st.rerun()
 
-# Display history
+# --- Message Rendering ---
 for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
+    # Handle the 'coder' role mapping to 'assistant' for UI consistency
+    role = "assistant" if msg["role"] == "coder" else msg["role"]
+    with st.chat_message(role):
         st.markdown(msg["content"])
 
-# 3. Chat Logic with your exact Template Code
-if prompt := st.chat_input("Message Qwen 2.5 32B..."):
+# --- Generation Logic ---
+if prompt := st.chat_input("Message to Qwen..."):
+    # Append User Message
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
+    # Generate Assistant Response
     with st.chat_message("assistant"):
-        # Setup Streamer
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        placeholder = st.empty()
+        full_response = ""
         
-        # YOUR EXACT LOGIC: Applying the chat template
+        # 1. Prepare Inputs
         inputs = tokenizer.apply_chat_template(
+            # Filter history to only include user/coder roles for the template
             st.session_state.messages,
             add_generation_prompt=True,
             tokenize=True,
             return_dict=True,
-            return_tensors="pt",
+            return_tensors="pt"
         ).to(model.device)
 
-        # Threading for live streaming
+        # 2. Setup Streamer
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+        # 3. Execution (Expert Note: use inference_mode for speed/memory)
         generation_kwargs = dict(
             **inputs,
             streamer=streamer,
-            max_new_tokens=512,
+            max_new_tokens=1024, # Increased for more robust answers
             do_sample=True,
             temperature=0.7,
+            top_p=0.9, # Added for higher quality sampling
             pad_token_id=tokenizer.eos_token_id
         )
-        
+
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        # Word-by-word UI update
-        placeholder = st.empty()
-        full_response = ""
+        # 4. Stream Handling
         for new_text in streamer:
             full_response += new_text
             placeholder.markdown(full_response + "▌")
-            
+        
         placeholder.markdown(full_response)
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
+        
+        # Store as 'coder' per original logic requirement
+        st.session_state.messages.append({"role": "coder", "content": full_response})