Update app.py

app.py

@@ -1,9 +1,8 @@
 from flask import Flask, request, jsonify
 import os
 import logging
-from transformers import pipeline
 
-# ====== Cache & Logging Setup ======
+# Set cache environment variables
 os.environ['HF_HOME'] = '/.cache/huggingface'
 os.environ['TRANSFORMERS_CACHE'] = '/.cache/huggingface/transformers'
 
@@ -13,57 +12,68 @@ logger = logging.getLogger(__name__)
 app = Flask(__name__)
 model = None
 
-# ====== Model Loader ======
 def load_model():
     global 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("Loading YOUR fine-tuned model...")
+        from transformers import pipeline
+        
+        model = pipeline(
+            "text-generation", 
+            model="kacperbb/phi-3.5-hf-finetuned",
+            trust_remote_code=True
+        )
+        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...")
         try:
-            logger.info(f"Loading {description}...")
             model = pipeline(
-                "text-generation",
-                model=model_name,
-                trust_remote_code=True,
-                cache_dir="/.cache/huggingface"
+                "text-generation", 
+                model="microsoft/Phi-3.5-mini-instruct",
+                trust_remote_code=True
             )
-            logger.info(f"✅ {description} loaded successfully!")
+            logger.info("✅ Base model loaded as fallback")
             return True
-        except Exception as e:
-            logger.error(f"❌ Failed to load {description}: {e}")
+        except Exception as e2:
+            logger.error(f"❌ Fallback failed: {e2}")
+            try:
+                model = pipeline("text-generation", model="gpt2")
+                logger.info("✅ GPT-2 fallback model loaded")
+                return True
+            except:
+                return False
 
-    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') or data.get('prompt', '')
+        prompt = data.get('inputs', data.get('prompt', ''))
         max_tokens = data.get('parameters', {}).get('max_new_tokens', 100)
-
+        
         if not prompt:
             return jsonify({"error": "No prompt provided"}), 400
-
-        if not model:
+        
+        if model:
+            result = model(prompt, max_new_tokens=max_tokens, do_sample=True)
+            response = result[0]['generated_text']
+        else:
             return jsonify({"error": "Model not loaded"}), 500
-
-        result = model(prompt, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
-        return jsonify([{"generated_text": result[0]['generated_text']}])
-
+        
+        return jsonify([{"generated_text": response}])
+    
     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():
@@ -76,7 +86,6 @@ def home():
         }
     })
 
-# ====== Start Server ======
 if __name__ == '__main__':
     logger.info("Starting Phi 3.5 API...")
     load_model()