Update app.py

app.py

@@ -1,67 +1,85 @@
 from flask import Flask, request, jsonify
-import requests
 import logging
-import os
+import torch
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
+model = None
+tokenizer = None
 
-# Your model endpoint
-MODEL_URL = "https://api-inference.huggingface.co/models/kacperbb/phi-3.5-merged-lora"
+def load_model():
+    global model, tokenizer
+    try:
+        logger.info("Loading YOUR fine-tuned model with transformers...")
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            "kacperbb/phi-3.5-merged-lora", 
+            trust_remote_code=True
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            "kacperbb/phi-3.5-merged-lora", 
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+            device_map="cpu"
+        )
+        logger.info("✅ YOUR fine-tuned model loaded successfully!")
+        return True
+    except Exception as e:
+        logger.error(f"❌ Error loading model: {e}")
+        return False
 
 @app.route('/generate', methods=['POST'])
 def generate():
+    global model, tokenizer
     try:
         data = request.json
         prompt = data.get('inputs', '')
-        max_tokens = data.get('parameters', {}).get('max_new_tokens', 100)
-        
-        # Get HF token from environment variable (you'll need to set this in Space settings)
-        hf_token = os.environ.get('HF_TOKEN')
+        max_tokens = data.get('parameters', {}).get('max_new_tokens', 50)
         
-        if hf_token:
-            # Call your actual model via HF Inference API
-            headers = {"Authorization": f"Bearer {hf_token}"}
-            payload = {
-                "inputs": prompt,
-                "parameters": {"max_new_tokens": max_tokens}
-            }
+        if model and tokenizer:
+            messages = [{"role": "user", "content": prompt}]
+            inputs = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_tensors="pt"
+            )
             
-            try:
-                response = requests.post(MODEL_URL, headers=headers, json=payload, timeout=30)
-                if response.status_code == 200:
-                    return jsonify(response.json())
-                else:
-                    logger.warning(f"HF API returned status {response.status_code}")
-            except requests.exceptions.RequestException as e:
-                logger.error(f"Request failed: {e}")
-        
-        # Fallback response if API call fails
-        return jsonify([{"generated_text": f"AI Response to '{prompt}': This would be generated by your fine-tuned Phi 3.5 model."}])
-        
+            with torch.no_grad():
+                outputs = model.generate(
+                    inputs, 
+                    max_new_tokens=max_tokens,
+                    do_sample=True,
+                    temperature=0.7
+                )
+            
+            response = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
+            return jsonify([{"generated_text": response}])
+        else:
+            return jsonify({"error": "Model not loaded"}), 500
+    
     except Exception as e:
-        logger.error(f"Error: {e}")
+        logger.error(f"Generation error: {e}")
         return jsonify({"error": str(e)}), 500
 
 @app.route('/health', methods=['GET'])
 def health():
-    has_token = bool(os.environ.get('HF_TOKEN'))
     return jsonify({
-        "status": "healthy", 
-        "model": "kacperbb/phi-3.5-merged-lora",
-        "has_token": has_token
+        "status": "healthy",
+        "model_loaded": model is not None
     })
 
 @app.route('/')
 def home():
     return jsonify({
-        "message": "Phi 3.5 Fine-tuned API Running!", 
-        "model": "kacperbb/phi-3.5-merged-lora",
-        "endpoint": "/generate"
+        "message": "Your Phi 3.5 Fine-tuned Model API",
+        "model": "kacperbb/phi-3.5-merged-lora"
     })
 
 if __name__ == '__main__':
-    logger.info("Starting Phi 3.5 Fine-tuned API...")
+    logger.info("Starting your fine-tuned model API...")
+    load_model()
     app.run(host='0.0.0.0', port=7860)