Update app.py

app.py

@@ -1,6 +1,7 @@
 from flask import Flask, request, jsonify
 import requests
 import os
+import json
 
 app = Flask(__name__)
 
@@ -19,22 +20,48 @@ def generate_text():
             "parameters": {
                 "max_new_tokens": 500,
                 "temperature": 0.7,
-                "top_p": 0.95
+                "top_p": 0.95,
+                "return_full_text": True
             }
         }
 
         # Make request to Hugging Face API
         response = requests.post(API_URL, headers=HEADERS, json=payload)
-        response.raise_for_status()  # Raise an exception for bad status codes
         
-        # Extract the generated text from response
-        generated_text = response.json()[0]['generated_text']
+        # Print response for debugging
+        print(f"Status Code: {response.status_code}")
+        print(f"Response Content: {response.text}")
+        
+        # Check if the model is still loading
+        if response.status_code == 503:
+            return jsonify({'error': 'Model is loading, please try again in a few seconds'}), 503
+            
+        # Check for other error status codes
+        response.raise_for_status()
+        
+        # Parse the response
+        response_data = response.json()
+        
+        # Handle different response formats
+        if isinstance(response_data, list) and len(response_data) > 0:
+            if 'generated_text' in response_data[0]:
+                generated_text = response_data[0]['generated_text']
+            else:
+                generated_text = response_data[0]
+        else:
+            generated_text = response_data
         
         return jsonify({'response': generated_text})
 
+    except requests.exceptions.RequestException as e:
+        print(f"API Request Error: {str(e)}")
+        return jsonify({'error': f'API Request Error: {str(e)}'}), 500
+    except json.JSONDecodeError as e:
+        print(f"JSON Decode Error: {str(e)}")
+        return jsonify({'error': 'Invalid response from API'}), 500
     except Exception as e:
-        print(f"Exception occurred: {e}")
-        return jsonify({'error': 'An unexpected error occurred.'}), 500
+        print(f"Unexpected Error: {str(e)}")
+        return jsonify({'error': f'An unexpected error occurred: {str(e)}'}), 500
 
 if __name__ == '__main__':
     app.run(debug=True)