Update app.py

app.py

@@ -1,26 +1,41 @@
 import os
 from flask import Flask, request, jsonify
-from dotenv import load_dotenv
-import requests
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import yaml
 
-load_dotenv()
+# Load environment variables
 api_key = os.getenv('HF_API_KEY')
 model_path = os.getenv('MODEL_PATH')
 
+# Initialize Flask app
 app = Flask(__name__)
 
-def get_model_predictions(text):
-    headers = {"Authorization": f"Bearer {api_key}"}
-    payload = {"inputs": text}
-    response = requests.post(f"https://api.huggingface.co/models/{model_path}", headers=headers, json=payload)
-    return response.json()
+# Load configuration
+with open('config.yaml', 'r') as file:
+    config = yaml.safe_load(file)
 
-@app.route('/predict', methods=['POST'])
-def predict():
+# Load the model and tokenizer
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(model_path)
+    return model, tokenizer
+
+model, tokenizer = load_model()
+
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs['input_ids'])
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+@app.route('/generate', methods=['POST'])
+def generate():
     data = request.get_json()
-    text = data['text']
-    prediction = get_model_predictions(text)
-    return jsonify(prediction)
+    prompt = data.get('prompt')
+    if prompt:
+        response_text = generate_text(prompt)
+        return jsonify({"response": response_text})
+    else:
+        return jsonify({"error": "No prompt provided"})
 
 if __name__ == '__main__':
     app.run(debug=True)