Update app.py

app.py

@@ -18,16 +18,15 @@ MODELS = {
         },
         "emoji": "🦁",
         "experimental": True,
-        "is_vision": False,  # Enable vision support for this model
-        "system_prompt_env": "ATLAS_FLASH_1215",  # Environment variable for system prompt
+        "is_vision": False,
+        "system_prompt_env": "ATLAS_FLASH_1215",
     },
 }
 
 # Profile pictures
-USER_PFP = "user.png"  # Hugging Face user avatar
-AI_PFP = "ai_pfp.png"  # Replace with the path to your AI's image or a URL
+USER_PFP = "user.png"
+AI_PFP = "ai_pfp.png"
 
-# Set page config (must be called only once and before any other Streamlit commands)
 st.set_page_config(
     page_title="Atlas Model Inference",
     page_icon="🦁 ",
@@ -39,15 +38,12 @@ st.set_page_config(
     }
 )
 
-# Custom CSS for blue sliders and button
 st.markdown(
     """
     <style>
-    /* Blue slider */
     .stSlider > div > div > div > div {
         background-color: #1f78b4 !important;
     }
-    /* Blue button */
     .stButton > button {
         background-color: #1f78b4 !important;
         color: white !important;
@@ -69,7 +65,6 @@ class AtlasInferenceApp:
             st.session_state.chat_history = []
 
     def clear_memory(self):
-        """Optimize memory management for CPU inference"""
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         gc.collect()
@@ -85,24 +80,22 @@ class AtlasInferenceApp:
 
             model_path = MODELS[model_key]["sizes"][model_size]
             
-            # Load Qwen-compatible tokenizer and model
             tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
             model = AutoModelForCausalLM.from_pretrained(
                 model_path,
-                device_map="auto",  # Use GPU if available
+                device_map="auto",
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 trust_remote_code=True,
                 low_cpu_mem_usage=True
             )
 
-            # Update session state
             st.session_state.current_model.update({
                 "tokenizer": tokenizer,
                 "model": model,
                 "config": {
                     "name": f"{MODELS[model_key]['name']} {model_size}",
                     "path": model_path,
-                    "system_prompt": os.getenv(MODELS[model_key]["system_prompt_env"], "Default system prompt"),  # Load system prompt from env
+                    "system_prompt": os.getenv(MODELS[model_key]["system_prompt_env"], "Default system prompt"),
                 }
             })
             return f"✅ {MODELS[model_key]['name']} {model_size} loaded successfully!"
@@ -114,18 +107,10 @@ class AtlasInferenceApp:
             return "⚠️ Please select and load a model first"
 
         try:
-            # Debugging: Check if config and system_prompt exist
-            if "config" not in st.session_state.current_model:
-                return "⚠️ Model configuration not found. Please load the model again."
-
-            system_prompt = st.session_state.current_model["config"].get("system_prompt", "Default system prompt")
+            system_prompt = st.session_state.current_model["config"]["system_prompt"]
             if not system_prompt:
-                system_prompt = "You are Atlas. You are developed by Spestly"  # Fallback if system_prompt is None
-
-            # Debugging: Print the system prompt for verification
-            st.write(f"System Prompt: {system_prompt}")
+                return "⚠️ System prompt not found for the selected model."
 
-            # Add the system instruction to guide the model's behavior
             prompt = f"{system_prompt}\n\n### Instruction:\n{message}\n\n### Response:"
 
             inputs = st.session_state.current_model["tokenizer"](
@@ -135,8 +120,6 @@ class AtlasInferenceApp:
                 truncation=True,
                 padding=True
             )
-
-            # Generate response without streaming
             with torch.no_grad():
                 output = st.session_state.current_model["model"].generate(
                     input_ids=inputs.input_ids,
@@ -151,7 +134,6 @@ class AtlasInferenceApp:
                 )
                 response = st.session_state.current_model["tokenizer"].decode(output[0], skip_special_tokens=True)
 
-            # Remove the prompt from the response
             if prompt in response:
                 response = response.replace(prompt, "").strip()
 
@@ -195,19 +177,16 @@ class AtlasInferenceApp:
 
         st.markdown("*⚠️ CAUTION: Atlas is an experimental model and this is just a preview. Responses may not be expected. Please double-check sensitive information!*")
 
-        # Display chat history
         for message in st.session_state.chat_history:
             with st.chat_message(
                 message["role"],
                 avatar=USER_PFP if message["role"] == "user" else AI_PFP
             ):
                 st.markdown(message["content"])
-                if "image" in message:
-                    st.image(message["image"], caption="Uploaded Image", use_container_width=True)  # Updated parameter
+                if "image" in message and message["image"]:
+                    st.image(message["image"], caption="Uploaded Image", use_column_width=True)
 
-        # Input box for user messages
         if prompt := st.chat_input("Message Atlas..."):
-            # Allow image upload if the model supports vision
             uploaded_image = None
             if MODELS[model_key]["is_vision"]:
                 uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
@@ -216,7 +195,7 @@ class AtlasInferenceApp:
             with st.chat_message("user", avatar=USER_PFP):
                 st.markdown(prompt)
                 if uploaded_image:
-                    st.image(uploaded_image, caption="Uploaded Image", use_container_width=True)  # Updated parameter
+                    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
 
             with st.chat_message("assistant", avatar=AI_PFP):
                 with st.spinner("Generating response..."):
@@ -233,4 +212,4 @@ def run():
         st.error(f"⚠️ Application Error: {str(e)}")
 
 if __name__ == "__main__":
-    run()
+    run()