asr24/StoryCraftAI

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