Update app.py

app.py

@@ -1,36 +1,24 @@
 import os
-from transformers import TFBertForSequenceClassification, BertTokenizerFast
+from flask import Flask, request, jsonify
+from dotenv import load_dotenv
+import tensorflow as tf
 
-def load_model(model_name):
-    try:
-        # Load TensorFlow model from Hugging Face
-        model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'))
-    except OSError:
-        # Fallback to PyTorch model if TensorFlow fails
-        model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'), from_pt=True)
-    return model
+load_dotenv()
+api_key = os.getenv('HF_API_KEY')
+model_path = os.getenv('MODEL_PATH')
 
-def load_tokenizer(model_name):
-    tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'))
-    return tokenizer
+app = Flask(__name__)
 
-def predict(text, model, tokenizer):
-    inputs = tokenizer(text, return_tensors="tf")
-    outputs = model(**inputs)
-    return outputs
+def load_model():
+    return tf.keras.models.load_model(model_path)
 
-def main():
-    model_name = os.getenv('MODEL_PATH')
-    if model_name is None:
-        raise ValueError("MODEL_PATH environment variable not set or is None")
+@app.route('/predict', methods=['POST'])
+def predict():
+    data = request.get_json()
+    text = data['text']
+    model = load_model()
+    prediction = model.predict([text])
+    return jsonify(prediction.tolist())
 
-    model = load_model(model_name)
-    tokenizer = load_tokenizer(model_name)
-
-    # Example prediction
-    text = "Sample input text"
-    result = predict(text, model, tokenizer)
-    print(result)
-
-if __name__ == "__main__":
-    main()
+if __name__ == '__main__':
+    app.run()