Spaces:
Running
Running
File size: 4,733 Bytes
3b74369 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
from flask import Flask, request, jsonify, render_template
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
import numpy as np
import pickle
from tensorflow.keras.preprocessing.text import Tokenizer
import sys
import tensorflow.keras.preprocessing.text as keras_text
sys.modules['keras.preprocessing.text'] = keras_text
app = Flask(__name__)
# Set the sequence length used during training (adjust as needed)
text_seq_len = 50
# Dictionaries to store models, histories, and tokenizers
models = {}
histories = {}
tokenizers = {}
@tf.keras.utils.register_keras_serializable()
class MyGRU(tf.keras.layers.GRU):
def __init__(self, *args, **kwargs):
kwargs.pop("time_major", None)
super().__init__(*args, **kwargs)
@tf.keras.utils.register_keras_serializable()
class MyLSTM(tf.keras.layers.LSTM):
def __init__(self, *args, **kwargs):
kwargs.pop("time_major", None)
super().__init__(*args, **kwargs)
@tf.keras.utils.register_keras_serializable()
class MyBI_LSTM(tf.keras.layers.Bidirectional):
def __init__(self, *args, **kwargs):
kwargs.pop("time_major", None)
super().__init__(*args, **kwargs)
@tf.keras.utils.register_keras_serializable()
class MyBI_GRU(tf.keras.layers.Bidirectional):
def __init__(self, *args, **kwargs):
kwargs.pop("time_major", None)
super().__init__(*args, **kwargs)
# Model 1: GRU (unchanged)
models['GRU'] = tf.keras.models.load_model("Models/GRU/GRU.h5", custom_objects={'GRU': MyGRU})
with open("Models/GRU/gru_history.pkl", "rb") as f:
histories['GRU'] = pickle.load(f)
with open("Models/GRU/gru_tokenizer.pkl", "rb") as f:
tokenizers['GRU'] = pickle.load(f)
# Model 2: LSTM – load with custom object mapping key 'LSTM'
models['LSTM'] = tf.keras.models.load_model("Models/LSTM/LSTM.h5", custom_objects={'LSTM': MyLSTM})
with open("Models/LSTM/lstm_history.pkl", "rb") as f:
histories['LSTM'] = pickle.load(f)
with open("Models/LSTM/lstm_tokenizer.pkl", "rb") as f:
tokenizers['LSTM'] = pickle.load(f)
# Model 3: Bidirectional-LSTM – map both 'Bidirectional' and 'LSTM'
models['Bidirectional-LSTM'] = tf.keras.models.load_model(
"Models/BIDIRECTIONAL_LSTM/Bi_di_LSTM.h5",
custom_objects={'Bidirectional': MyBI_LSTM, 'LSTM': MyLSTM}
)
with open("Models/BIDIRECTIONAL_LSTM/bi_di_lstm_history.pkl", "rb") as f:
histories['Bidirectional-LSTM'] = pickle.load(f)
with open("Models/BIDIRECTIONAL_LSTM/bi_di_lstm_tokenizer.pkl", "rb") as f:
tokenizers['Bidirectional-LSTM'] = pickle.load(f)
# Model 4: Bidirectional-GRU – map both 'Bidirectional' and 'GRU'
models['Bidirectional-GRU'] = tf.keras.models.load_model(
"Models/BIDIRECTIONAL_GRU/Bi_di_GRU.h5",
custom_objects={'Bidirectional': MyBI_GRU, 'GRU': MyGRU}
)
with open("Models/BIDIRECTIONAL_GRU/bi_di_gru_history.pkl", "rb") as f:
histories['Bidirectional-GRU'] = pickle.load(f)
with open("Models/BIDIRECTIONAL_GRU/bi_di_gru_tokenizer.pkl", "rb") as f:
tokenizers['Bidirectional-GRU'] = pickle.load(f)
def generate_story(model, tokenizer, text_seq_len, seed_text, n_words):
text = []
for _ in range(n_words):
# Convert seed text to a sequence of integers
encoded = tokenizer.texts_to_sequences([seed_text])[0]
# Pad the sequence to ensure consistent length
encoded = pad_sequences([encoded], maxlen=text_seq_len, padding='pre')
# Predict the next word probabilities
pred = model.predict(encoded)
# Choose the word with the highest probability
y_pred = np.argmax(pred, axis=1)[0]
predicted_word = ''
# Map the integer back to the word
for word, index in tokenizer.word_index.items():
if index == y_pred:
predicted_word = word
break
seed_text = seed_text + ' ' + predicted_word
text.append(predicted_word)
return ' '.join(text)
@app.route("/")
def home():
return render_template("index.html", histories=histories)
@app.route("/generate-story", methods=["POST"])
def generate_story_endpoint():
data = request.get_json()
seed_text = data.get("seed_text", "")
num_words = data.get("num_words", 100)
model_name = data.get("model_name", "GRU")
selected_model = models.get(model_name)
selected_tokenizer = tokenizers.get(model_name)
if selected_model is None or selected_tokenizer is None:
return jsonify({"error": "Model or tokenizer not found"}), 400
generated_story = generate_story(selected_model, selected_tokenizer, text_seq_len, seed_text, num_words)
return jsonify({"generated_story": generated_story})
if __name__ == "__main__":
app.run(debug=True, port=7860)
|