Upload Flask app

app.py

@@ -0,0 +1,172 @@
+from flask import Flask, render_template, request, jsonify, stream_template
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import re
+import gc
+import json
+import threading
+import time
+
+app = Flask(__name__)
+
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
+model_loaded = False
+loading_status = "Initializing..."
+
+def load_model():
+    """Load the model and tokenizer optimized for CPU"""
+    global model, tokenizer, model_loaded, loading_status
+    
+    try:
+        loading_status = "Loading tokenizer..."
+        print("Loading AEGIS Conduct Economic Analysis Model for CPU...")
+        
+        # Load tokenizer first
+        tokenizer = AutoTokenizer.from_pretrained(
+            "Gaston895/aegisconduct",
+            trust_remote_code=True
+        )
+        
+        loading_status = "Loading model (this may take a few minutes)..."
+        
+        # Load model optimized for CPU
+        model = AutoModelForCausalLM.from_pretrained(
+            "Gaston895/aegisconduct",
+            torch_dtype=torch.float16,
+            device_map="cpu",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        # Force garbage collection
+        gc.collect()
+        
+        loading_status = "Model loaded successfully!"
+        print("Model loaded successfully on CPU!")
+        model_loaded = True
+        return True
+        
+    except Exception as e:
+        loading_status = f"Error loading model: {str(e)}"
+        print(f"Error loading model: {e}")
+        model_loaded = False
+        return False
+
+def format_response(text):
+    """Clean and format the model response"""
+    # Remove thinking tags if present
+    text = re.sub(r'<thinking>.*?</thinking>', '', text, flags=re.DOTALL)
+    
+    # Clean up extra whitespace
+    text = re.sub(r'\n\s*\n', '\n\n', text)
+    text = text.strip()
+    
+    return text
+
+def generate_response(message, history=None, temperature=0.7, max_tokens=128):
+    """Generate response from the model optimized for CPU"""
+    global model, tokenizer, model_loaded
+    
+    if not model_loaded or model is None or tokenizer is None:
+        return "Model is still loading... Please wait a moment and try again."
+    
+    try:
+        # Build conversation context
+        conversation = ""
+        if history:
+            # Only use last 2 exchanges to save memory
+            recent_history = history[-2:] if len(history) > 2 else history
+            for exchange in recent_history:
+                conversation += f"User: {exchange['user']}\nAssistant: {exchange['assistant']}\n\n"
+        
+        # Add current message
+        conversation += f"User: {message}\nAssistant:"
+        
+        # Tokenize input with strict length limit for CPU
+        inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=512)
+        
+        # Generate response with CPU-optimized settings
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=True,
+                top_p=0.9,
+                top_k=50,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                use_cache=True,
+                num_beams=1
+            )
+        
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract only the new response
+        response = response[len(conversation):].strip()
+        
+        # Format and clean response
+        response = format_response(response)
+        
+        # Clean up memory after generation
+        gc.collect()
+        
+        return response if response else "I apologize, but I couldn't generate a proper response. Please try rephrasing your question."
+        
+    except Exception as e:
+        return f"Error generating response: {str(e)}. Please try a shorter question."
+
+@app.route('/')
+def index():
+    """Main chat interface"""
+    return render_template('index.html')
+
+@app.route('/status')
+def status():
+    """Get model loading status"""
+    return jsonify({
+        'loaded': model_loaded,
+        'status': loading_status
+    })
+
+@app.route('/chat', methods=['POST'])
+def chat():
+    """Handle chat messages"""
+    data = request.json
+    message = data.get('message', '').strip()
+    history = data.get('history', [])
+    temperature = float(data.get('temperature', 0.7))
+    max_tokens = int(data.get('max_tokens', 128))
+    
+    if not message:
+        return jsonify({'error': 'No message provided'}), 400
+    
+    # Generate response
+    response = generate_response(message, history, temperature, max_tokens)
+    
+    return jsonify({
+        'response': response,
+        'timestamp': time.time()
+    })
+
+@app.route('/clear', methods=['POST'])
+def clear_chat():
+    """Clear chat history and free memory"""
+    gc.collect()
+    return jsonify({'status': 'cleared'})
+
+# Start model loading in background thread
+def start_model_loading():
+    load_model()
+
+if __name__ == '__main__':
+    # Start model loading in background
+    loading_thread = threading.Thread(target=start_model_loading)
+    loading_thread.daemon = True
+    loading_thread.start()
+    
+    app.run(host='0.0.0.0', port=7860, debug=False)