Aifybox / app.py
shing-dev's picture
changed app.py a bit for hf
2ccff21
Raw History Blame Contribute Delete
2.22 kB
from flask import Flask, request, jsonify
from flask_cors import CORS
import tensorflow as tf
import pickle
import numpy as np
from tensorflow.keras.preprocessing.sequence import pad_sequences
import os
# import requests
import gdown
app = Flask(__name__)
CORS(app)
# Load tokenizer and model
with open("tokenizer.pkl", "rb") as f:
tokenizer = pickle.load(f)
MODEL_PATH = "model.keras"
MODEL_URL = (
"https://drive.google.com/uc?export=download&id=1uqqZiZsmI2wnxG6r9Z4IQq94rT9fGwgH"
)
def download_model():
if not os.path.exists(MODEL_PATH):
print("Downloading model using gdown...")
# Use gdown to download the file
gdown.download(MODEL_URL, MODEL_PATH, quiet=False)
if os.path.exists(MODEL_PATH):
print("Model downloaded successfully.")
else:
print("Error: Model download failed.")
download_model()
print("Loading model...")
model = tf.keras.models.load_model(MODEL_PATH)
MAX_LEN = 30
def predict_next_words(text, top_k=3):
# Tokenize input text
sequence = tokenizer.texts_to_sequences([text])[0]
if len(sequence) == 0: # If tokenization results in empty list
return []
# Keep only the last 30 tokens if input is too long
sequence = sequence[-MAX_LEN:]
# Pad sequence to required length (post-padding)
sequence = pad_sequences([sequence], maxlen=MAX_LEN, padding="post")
# Predict next word probabilities
predictions = model.predict(sequence)[0]
# Get top-k word indices
top_indices = np.argsort(predictions)[-top_k:][::-1]
# Convert indices back to words
word_index = tokenizer.index_word
top_words = [word_index.get(i, "") for i in top_indices]
return [word for word in top_words if word]
@app.route("/predict", methods=["POST"])
def predict():
data = request.get_json()
text = data.get("text", "").strip()
if not text:
return jsonify({"error": "No text provided"}), 400
predictions = predict_next_words(text)
return jsonify({"predictions": predictions})
@app.route("/")
def home():
return "Your Flask app is running on Hugging Face Spaces!"
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)