Update app.py

app.py

@@ -1,67 +1,47 @@
 from flask import Flask, request, jsonify
-import requests
+from huggingface_hub import InferenceClient
 import os
-import json
 
 app = Flask(__name__)
 
-API_URL = "https://api-inference.huggingface.co/models/deepseek-ai/deepseek-coder-6.7b-base"
-HEADERS = {"Authorization": f"Bearer {os.getenv('HUGGING_FACE_API_KEY')}"}
+# Initialize the InferenceClient
+client = InferenceClient(
+    token=os.getenv('HUGGING_FACE_API_KEY')  # Make sure to set this environment variable
+)
 
 @app.route('/', methods=['GET'])
 def generate_text():
     try:
+        # Verify API key is set
+        if not os.getenv('HUGGING_FACE_API_KEY'):
+            return jsonify({'error': 'HUGGING_FACE_API_KEY environment variable is not set'}), 500
+
         # Get the user's prompt from query parameters
         user_prompt = request.args.get('prompt', 'What is the capital of France?')
 
-        # Prepare the payload
-        payload = {
-            "inputs": user_prompt,
-            "parameters": {
-                "max_new_tokens": 500,
-                "temperature": 0.7,
-                "top_p": 0.95,
-                "return_full_text": True
+        # Prepare the messages
+        messages = [
+            {
+                "role": "user",
+                "content": user_prompt
             }
-        }
+        ]
 
-        # Make request to Hugging Face API
-        response = requests.post(API_URL, headers=HEADERS, json=payload)
-        
-        # 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
+        # Generate completion
+        completion = client.chat.completions.create(
+            model="deepseek-ai/DeepSeek-R1",
+            messages=messages,
+            max_tokens=500
+        )
+
+        # Extract the response
+        response_text = completion.choices[0].message.content
         
-        return jsonify({'response': generated_text})
+        return jsonify({'response': response_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"Unexpected Error: {str(e)}")
-        return jsonify({'error': f'An unexpected error occurred: {str(e)}'}), 500
+        print(f"Error occurred: {str(e)}")
+        return jsonify({'error': f'An error occurred: {str(e)}'}), 500
 
 if __name__ == '__main__':
     app.run(debug=True)