Update app.py

app.py

@@ -1,8 +1,9 @@
 from flask import Flask, request, jsonify
 import os
 import logging
+from transformers import pipeline
 
-# Set cache environment variables
+# ====== Cache & Logging Setup ======
 os.environ['HF_HOME'] = '/.cache/huggingface'
 os.environ['TRANSFORMERS_CACHE'] = '/.cache/huggingface/transformers'
 
@@ -12,72 +13,57 @@ logger = logging.getLogger(__name__)
 app = Flask(__name__)
 model = None
 
+# ====== Model Loader ======
 def load_model():
     global model
-    try:
-        logger.info("Loading YOUR fine-tuned model...")
-        from transformers import pipeline
-        
-        # Try to use your model with cache settings
-        model = pipeline(
-            "text-generation", 
-            model="kacperbb/phi-3.5-hf-finetuned",
-            trust_remote_code=True,
-            cache_dir="/.cache/huggingface"
-        )
-        logger.info("✅ YOUR fine-tuned model loaded successfully!")
-        return True
-    except Exception as e:
-        logger.error(f"❌ Error loading your model: {e}")
-        logger.info("Trying with base model...")
+    models_to_try = [
+        ("kacperbb/phi-3.5-hf-finetuned", "YOUR fine-tuned model"),
+        ("microsoft/Phi-3.5-mini-instruct", "Base model"),
+        ("microsoft/DialoGPT-medium", "Smaller fallback model")
+    ]
+
+    for model_name, description in models_to_try:
         try:
+            logger.info(f"Loading {description}...")
             model = pipeline(
-                "text-generation", 
-                model="microsoft/Phi-3.5-mini-instruct",
+                "text-generation",
+                model=model_name,
                 trust_remote_code=True,
                 cache_dir="/.cache/huggingface"
             )
-            logger.info("✅ Base model loaded as fallback")
+            logger.info(f"✅ {description} loaded successfully!")
             return True
-        except Exception as e2:
-            logger.error(f"❌ Fallback failed: {e2}")
-            # Last resort - use a smaller model
-            try:
-                model = pipeline("text-generation", model="microsoft/DialoGPT-medium")
-                logger.info("✅ Smaller fallback model loaded")
-                return True
-            except:
-                return False
+        except Exception as e:
+            logger.error(f"❌ Failed to load {description}: {e}")
 
+    logger.error("❌ All model loading attempts failed!")
+    return False
+
+# ====== Endpoints ======
 @app.route('/generate', methods=['POST'])
 def generate_text():
     global model
     try:
         data = request.json
-        prompt = data.get('inputs', data.get('prompt', ''))
+        prompt = data.get('inputs') or data.get('prompt', '')
         max_tokens = data.get('parameters', {}).get('max_new_tokens', 100)
-        
+
         if not prompt:
             return jsonify({"error": "No prompt provided"}), 400
-        
-        if model:
-            result = model(prompt, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
-            response = result[0]['generated_text']
-        else:
+
+        if not model:
             return jsonify({"error": "Model not loaded"}), 500
-        
-        return jsonify([{"generated_text": response}])
-    
+
+        result = model(prompt, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
+        return jsonify([{"generated_text": result[0]['generated_text']}])
+
     except Exception as e:
         logger.error(f"Generation error: {e}")
         return jsonify({"error": str(e)}), 500
 
 @app.route('/health', methods=['GET'])
 def health():
-    return jsonify({
-        "status": "healthy",
-        "model_loaded": model is not None
-    })
+    return jsonify({"status": "healthy", "model_loaded": model is not None})
 
 @app.route('/', methods=['GET'])
 def home():
@@ -90,6 +76,7 @@ def home():
         }
     })
 
+# ====== Start Server ======
 if __name__ == '__main__':
     logger.info("Starting Phi 3.5 API...")
     load_model()