Upload app.py

app.py

@@ -53,6 +53,9 @@ HTML_TEMPLATE = """
         <div id="chat-container" class="chat-container">
             <div class="message ai-message">
                 Hello! I'm AEGIS Economics AI. Ask me about economic policies, market analysis, or financial strategies.
+                <div id="model-status" style="font-size: 0.8em; color: #666; margin-top: 5px;">
+                    Checking model status...
+                </div>
             </div>
         </div>
         
@@ -63,6 +66,49 @@ HTML_TEMPLATE = """
     </div>
 
     <script>
+        // Check model status on page load
+        async function checkModelStatus() {
+            try {
+                const response = await fetch('/health');
+                const data = await response.json();
+                const statusDiv = document.getElementById('model-status');
+                
+                if (data.model_loaded) {
+                    statusDiv.textContent = '✅ Model loaded and ready!';
+                    statusDiv.style.color = '#28a745';
+                } else {
+                    statusDiv.textContent = '⏳ Model loading... Please wait.';
+                    statusDiv.style.color = '#ffc107';
+                    // Try to load model
+                    setTimeout(tryLoadModel, 2000);
+                }
+            } catch (error) {
+                const statusDiv = document.getElementById('model-status');
+                statusDiv.textContent = '❌ Connection error';
+                statusDiv.style.color = '#dc3545';
+            }
+        }
+
+        async function tryLoadModel() {
+            try {
+                const response = await fetch('/load_model', { method: 'POST' });
+                const data = await response.json();
+                
+                if (data.success) {
+                    const statusDiv = document.getElementById('model-status');
+                    statusDiv.textContent = '✅ Model loaded successfully!';
+                    statusDiv.style.color = '#28a745';
+                } else {
+                    setTimeout(checkModelStatus, 5000); // Check again in 5 seconds
+                }
+            } catch (error) {
+                setTimeout(checkModelStatus, 5000);
+            }
+        }
+
+        // Call on page load
+        window.onload = checkModelStatus;
+
         function handleKeyPress(event) {
             if (event.key === 'Enter') {
                 sendMessage();
@@ -152,8 +198,8 @@ def load_model():
         logger.info(f"Loading model from {model_repo}...")
         model = AutoModelForCausalLM.from_pretrained(
             model_repo,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
+            torch_dtype=torch.float16,  # Changed from bfloat16 for better compatibility
+            device_map="cpu",           # Force CPU for HF Spaces compatibility
             trust_remote_code=True,
             use_auth_token=False,
             low_cpu_mem_usage=True
@@ -164,13 +210,31 @@ def load_model():
         
     except Exception as e:
         logger.error(f"Error loading model from HF: {str(e)}")
-        return False
+        # Try alternative loading method
+        try:
+            logger.info("Trying alternative loading method...")
+            tokenizer = AutoTokenizer.from_pretrained(
+                "Qwen/Qwen2-1.5B",  # Fallback to base model
+                trust_remote_code=True
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                "Qwen/Qwen2-1.5B",
+                torch_dtype=torch.float16,
+                device_map="cpu",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+            logger.info("Fallback model loaded successfully!")
+            return True
+        except Exception as e2:
+            logger.error(f"Fallback loading also failed: {str(e2)}")
+            return False
 
 def generate_response(prompt):
     """Generate response using the loaded model"""
     try:
         if model is None or tokenizer is None:
-            return "Model not loaded. Please wait..."
+            return "Model is still loading, please wait a moment and try again..."
         
         # Economics-focused system prompt
         system_prompt = """You are AEGIS Economics AI, an expert economic analyst and policy advisor. 
@@ -180,17 +244,18 @@ def generate_response(prompt):
         full_prompt = f"{system_prompt}\n\nUser: {prompt}\nAssistant:"
         
         # Tokenize input
-        inputs = tokenizer(full_prompt, return_tensors="pt")
+        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=1024)
         
         # Generate response
         with torch.no_grad():
             outputs = model.generate(
                 inputs.input_ids,
-                max_new_tokens=512,
+                max_new_tokens=256,  # Reduced for faster generation
                 temperature=0.7,
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
-                repetition_penalty=1.1
+                repetition_penalty=1.1,
+                no_repeat_ngram_size=3
             )
         
         # Decode response
@@ -204,7 +269,7 @@ def generate_response(prompt):
         
     except Exception as e:
         logger.error(f"Error generating response: {str(e)}")
-        return "I apologize, but I'm having trouble processing your request right now."
+        return "I apologize, but I'm having trouble processing your request right now. Please try again in a moment."
 
 @app.route('/')
 def home():
@@ -236,16 +301,34 @@ def health():
     return jsonify({
         'status': 'healthy',
         'model_loaded': model is not None,
-        'tokenizer_loaded': tokenizer is not None
+        'tokenizer_loaded': tokenizer is not None,
+        'model_info': 'Gaston895/Aegisecon1' if model is not None else 'Not loaded'
     })
 
+@app.route('/load_model', methods=['POST'])
+def load_model_endpoint():
+    """Endpoint to trigger model loading"""
+    try:
+        success = load_model()
+        return jsonify({
+            'success': success,
+            'model_loaded': model is not None,
+            'tokenizer_loaded': tokenizer is not None
+        })
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
 if __name__ == '__main__':
     # Load model on startup
     logger.info("Starting AEGIS Economics AI...")
     
-    if load_model():
+    # Try to load model, but don't fail if it doesn't work
+    logger.info("Attempting to load model...")
+    model_loaded = load_model()
+    
+    if model_loaded:
         logger.info("Model loaded successfully, starting server...")
-        app.run(host='0.0.0.0', port=7860, debug=False)
     else:
-        logger.error("Failed to load model, exiting...")
-        exit(1)
+        logger.warning("Model failed to load, starting server anyway. Model can be loaded via /load_model endpoint.")
+    
+    app.run(host='0.0.0.0', port=7860, debug=False)