Update app.py

app.py

@@ -11,43 +11,45 @@ logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
 model = None
+tokenizer = None
 
 def load_model():
-    global model
+    global model, tokenizer
     try:
         logger.info("Loading YOUR fine-tuned model...")
-        from transformers import pipeline
+        from transformers import AutoTokenizer, AutoModelForCausalLM
         
-        model = pipeline(
-            "text-generation", 
-            model="kacperbb/phi-3.5-merged-lora",
+        # Load model and tokenizer separately for better control
+        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="auto",
+            device_map="cpu"
+        )
+        
+        # Set pad token if not set
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
         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:
-            model = pipeline(
-                "text-generation", 
-                model="microsoft/Phi-3.5-mini-instruct",
-                trust_remote_code=True
-            )
-            logger.info("✅ Base model loaded as fallback")
+            from transformers import pipeline
+            model = pipeline("text-generation", model="gpt2")
+            logger.info("✅ Fallback model loaded")
             return True
-        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
+        except:
+            return False
 
 @app.route('/generate', methods=['POST'])
 def generate_text():
-    global model
+    global model, tokenizer
     try:
         data = request.json
         prompt = data.get('inputs', data.get('prompt', ''))
@@ -56,11 +58,27 @@ def generate_text():
         if not prompt:
             return jsonify({"error": "No prompt provided"}), 400
         
-        if model:
+        if model and tokenizer and hasattr(model, 'generate'):
+            # Use model directly
+            inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    inputs.input_ids,
+                    attention_mask=inputs.attention_mask,
+                    max_new_tokens=max_tokens,
+                    do_sample=True,
+                    temperature=0.7,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+            
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        elif model and hasattr(model, '__call__'):
+            # Use pipeline
             result = model(prompt, max_new_tokens=max_tokens, do_sample=True)
             response = result[0]['generated_text']
         else:
-            return jsonify({"error": "Model not loaded"}), 500
+            return jsonify({"error": "Model not properly loaded"}), 500
         
         return jsonify([{"generated_text": response}])
     
@@ -87,6 +105,7 @@ def home():
     })
 
 if __name__ == '__main__':
+    import torch
     logger.info("Starting Phi 3.5 API...")
     load_model()
     port = int(os.environ.get('PORT', 7860))