Update app.py

app.py

@@ -2,139 +2,70 @@ from flask import Flask, request, jsonify
 from flask_cors import CORS
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import os
 import time
 
 app = Flask(__name__)
 CORS(app)
 
-# Global variables for model caching
 model = None
 tokenizer = None
 model_loaded = False
 
-# Simba system message
 SIMBA_SYSTEM = """You are Simba from The Lion King. You're brave, playful, and wise. 
 Speak with royal confidence but also warmth and humor. Remember: "Hakuna Matata", 
 relationships with Nala, Timon, Pumbaa, and your journey to reclaim Pride Rock.
 Keep responses under 2 sentences and stay in character."""
 
 def load_model():
-    """Load model with optimizations - called once at startup"""
     global model, tokenizer, model_loaded
-    
     if model_loaded:
         return
     
-    print("🚀 Loading optimized Qwen2.5-0.5B model...")
-    start_time = time.time()
-    
+    print("Loading Qwen2.5-0.5B model...")
     model_name = "Qwen/Qwen2.5-0.5B-Instruct"
     
-    # 🎯 SPEED OPTIMIZATION 1: Use bfloat16 for faster inference
-    torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    
-    # 🎯 SPEED OPTIMIZATION 2: Faster tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        trust_remote_code=True,
-        padding_side="left"  # Better for batch processing
-    )
-    
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     
-    # 🎯 SPEED OPTIMIZATION 3: Optimized model loading
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch_dtype,
+        torch_dtype=torch.float16,
         device_map="auto",
-        trust_remote_code=True,
-        attn_implementation="sdpa",  # Flash Attention 2 for speed
-        use_cache=True,  # Faster generation
-        low_cpu_mem_usage=True,
+        trust_remote_code=True
     )
-    
-    # 🎯 SPEED OPTIMIZATION 4: Compile model for faster inference (PyTorch 2.0+)
-    if hasattr(torch, 'compile') and torch.cuda.is_available():
-        print("🔧 Compiling model for maximum speed...")
-        model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
-    
     model_loaded = True
-    load_time = time.time() - start_time
-    print(f"✅ Model loaded in {load_time:.2f} seconds!")
+    print("Model loaded!")
 
-# Load model when app starts
 load_model()
 
 def generate_response(user_message):
-    """Generate optimized response with speed enhancements"""
-    start_time = time.time()
-    
-    # Create conversation format for Qwen
     messages = [
         {"role": "system", "content": SIMBA_SYSTEM},
         {"role": "user", "content": user_message}
     ]
     
-    # 🎯 SPEED OPTIMIZATION 5: Efficient template application
-    text = tokenizer.apply_chat_template(
-        messages, 
-        tokenize=False, 
-        add_generation_prompt=True
-    )
-    
-    # 🎯 SPEED OPTIMIZATION 6: Optimized tokenization
-    inputs = tokenizer(
-        text, 
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=512
-    ).to(model.device)
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
     
-    # 🎯 SPEED OPTIMIZATION 7: Faster generation parameters
     with torch.no_grad():
-        # Use inference mode for speed
-        with torch.inference_mode():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=100,  # Reduced for speed
-                temperature=0.7,
-                do_sample=True,
-                top_p=0.9,
-                top_k=40,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                num_return_sequences=1,
-                early_stopping=True
-            )
-    
-    # 🎯 SPEED OPTIMIZATION 8: Efficient decoding
-    response = tokenizer.decode(
-        outputs[0][inputs['input_ids'].shape[1]:], 
-        skip_special_tokens=True
-    )
-    
-    generation_time = time.time() - start_time
-    print(f"⚡ Response generated in {generation_time:.2f} seconds")
-    
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
     return response.strip()
 
 @app.route('/')
 def home():
-    return jsonify({
-        "message": "Simba AI API is running! 🦁", 
-        "status": "optimized",
-        "model": "Qwen2.5-0.5B-Instruct"
-    })
+    return jsonify({"message": "Simba AI API is running! 🦁"})
 
-@app.route('/api/chat', methods=['POST', 'OPTIONS'])
+@app.route('/api/chat', methods=['POST'])
 def chat():
-    if request.method == 'OPTIONS':
-        return '', 200
-        
     try:
         data = request.get_json()
         user_message = data.get('message', '')
@@ -142,54 +73,18 @@ def chat():
         if not user_message:
             return jsonify({"error": "No message provided"}), 400
         
-        # 🎯 SPEED OPTIMIZATION 9: Input validation and truncation
-        if len(user_message) > 500:
-            user_message = user_message[:500] + "..."
-        
         response = generate_response(user_message)
         
         return jsonify({
             "response": response,
-            "status": "success",
-            "model": "Qwen2.5-0.5B"
+            "status": "success"
         })
         
     except Exception as e:
-        print(f"❌ Error: {str(e)}")
         return jsonify({
-            "error": "Hakuna Matata! Even kings have technical issues. Try again!",
+            "error": str(e),
             "status": "error"
         }), 500
 
-@app.route('/health')
-def health():
-    return jsonify({
-        "status": "healthy", 
-        "model_loaded": model_loaded,
-        "device": str(model.device) if model else "none"
-    })
-
-@app.route('/status')
-def status():
-    gpu_info = "CPU"
-    if torch.cuda.is_available():
-        gpu_info = f"GPU: {torch.cuda.get_device_name()}, Memory: {torch.cuda.memory_allocated()/1024**3:.2f}GB"
-    
-    return jsonify({
-        "status": "running",
-        "model": "Qwen2.5-0.5B-Instruct",
-        "optimizations": "enabled",
-        "hardware": gpu_info,
-        "torch_version": torch.__version__
-    })
-
-# CORS headers
-@app.after_request
-def after_request(response):
-    response.headers.add('Access-Control-Allow-Origin', '*')
-    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
-    response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
-    return response
-
 if __name__ == '__main__':
-    app.run(debug=False, host='0.0.0.0', port=7860)
+    app.run(debug=True, host='0.0.0.0', port=7860)