asr24/StoryCraftAI

1#!/usr/bin/env python32"""3Automatic Story Generation Script4 5This script scrapes short stories from http://www.classicshorts.com/,6preprocesses the text, builds and trains a neural network model to predict the next word,7and generates new story text from a seed phrase.8 9Available model types:10    - "bi_di_gru"   : Bidirectional GRU (default)11    - "bi_di_lstm"  : Bidirectional LSTM12    - "gru"         : Unidirectional GRU13    - "lstm"        : Unidirectional LSTM14"""15 16import os17import re18import string19import pickle20import requests21import nltk22from bs4 import BeautifulSoup23import numpy as np24import matplotlib.pyplot as plt25import tensorflow as tf26 27from nltk.tokenize import sent_tokenize, word_tokenize28from tensorflow.keras.preprocessing.text import Tokenizer29from tensorflow.keras.preprocessing.sequence import pad_sequences30from tensorflow.keras.utils import to_categorical31from tensorflow.keras.models import Sequential32from tensorflow.keras.layers import Dense, Dropout, GRU, LSTM, Embedding, Bidirectional33 34# Download necessary NLTK data35nltk.download('punkt_tab')36 37gpus = tf.config.list_physical_devices('GPU')38if gpus:39    try:40        for gpu in gpus:41            tf.config.experimental.set_memory_growth(gpu, True)42        print("Using GPU(s):", tf.config.list_physical_devices('GPU'))43    except RuntimeError as e:44        print("Error setting GPU memory growth:", e)45else:46    print("No GPU detected. Check your runtime settings.")47 48os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"49print("TensorFlow Version:", tf.__version__)50 51def get_page(url: str, headers: dict) -> BeautifulSoup:52    """53    Retrieve and parse a webpage.54 55    Args:56        url (str): URL of the page.57        headers (dict): HTTP headers for the request.58 59    Returns:60        BeautifulSoup: Parsed HTML content.61    """62    response = requests.get(url, headers=headers)63    try:64        response.raise_for_status()65    except Exception:66        pass  # In production, log this error67    return BeautifulSoup(response.text, 'html.parser')68 69 70def get_story_text(link: str) -> str:71    """72    Fetch the full story text from a given URL by cleaning the HTML.73 74    Args:75        link (str): Story URL.76 77    Returns:78        str: Cleaned story text.79    """80    regex = re.compile(r'[\n\r\t]')81    headers = {'User-Agent': 'Mozilla/5.0'}82    page_html = get_page(link, headers)83    paragraphs = page_html.find_all("div", class_="StoryPara")84    total_text = ""85    for paragraph in paragraphs:86        total_text += regex.sub(" ", paragraph.text.strip())87    return total_text88 89 90def get_listings(max_stories: int = 100) -> str:91    """92    Scrape the Classic Short Stories listings and return concatenated story text.93 94    Args:95        max_stories (int): Maximum number of stories to scrape.96 97    Returns:98        str: Combined text of all scraped stories.99    """100    story_count = 0101    bad_titles = {'tlrm', 'fiddler', 'frog', 'ItalianMaster', 'luck'}102    base_url = "http://www.classicshorts.com"103    listings_url = f"{base_url}/bib.html"104    headers = {'User-Agent': 'Mozilla/5.0'}105    raw_text = ""106    107    page_html = get_page(listings_url, headers)108    listing_elements = page_html.find_all("div", class_="biolisting")109    110    for elem in listing_elements:111        story_id = elem.attrs['onclick'][11:-2]112        if story_id not in bad_titles:113            current_url = f"{base_url}/stories/{story_id}.html"114            raw_text += get_story_text(current_url)115            story_count += 1116        if story_count == max_stories:117            break118    return raw_text119 120def clean_text(sentences: list) -> list:121    """122    Tokenize and clean sentences by removing punctuation, non-alphabetic tokens,123    and converting text to lowercase.124 125    Args:126        sentences (list): List of sentence strings.127 128    Returns:129        list: List of cleaned tokens.130    """131    tokens = []132    for sentence in sentences:133        tokens.extend(word_tokenize(sentence))134    translator = str.maketrans('', '', string.punctuation)135    tokens = [token.translate(translator) for token in tokens]136    tokens = [token.lower() for token in tokens if token.isalpha()]137    return tokens138 139 140def prepare_data() -> tuple:141    """142    Scrape, tokenize, and generate training sequences from the text data.143 144    Returns:145        tuple: (X, y, vocabulary_size, seq_length, tokenizer)146            X (np.ndarray): Input sequences.147            y (np.ndarray): One-hot encoded target words.148            vocabulary_size (int): Vocabulary size.149            seq_length (int): Length of input sequences.150            tokenizer (Tokenizer): Fitted tokenizer.151    """152    # Scrape and pre-process stories.153    stories = get_listings()154    # Optionally remove header noise (adjust slicing if needed)155    stories = stories[81:]156    sentences = sent_tokenize(stories)157    tokens = clean_text(sentences)158    159    # Create sequences of 51 tokens (50 as input, 1 as output)160    seq_total_length = 50 + 1161    lines = []162    for i in range(seq_total_length, len(tokens)):163        sequence = tokens[i - seq_total_length:i]164        lines.append(' '.join(sequence))165        if i > 120000:  # Limit data size for faster training166            break167 168    # Tokenize the sequences into integers.169    tokenizer = Tokenizer()170    tokenizer.fit_on_texts(lines)171    sequences = tokenizer.texts_to_sequences(lines)172    sequences = np.array(sequences)173    174    # Split sequences into inputs (X) and targets (y)175    X, y = sequences[:, :-1], sequences[:, -1]176    vocabulary_size = len(tokenizer.word_index) + 1177    y = to_categorical(y, num_classes=vocabulary_size)178    seq_length = X.shape[1]179    180    return X, y, vocabulary_size, seq_length, tokenizer181 182def generate_story(model: tf.keras.Model,183                   tokenizer: Tokenizer,184                   text_seq_len: int,185                   seed_text: str,186                   n_words: int) -> str:187    """188    Generate new text based on a seed phrase by predicting one word at a time.189 190    Args:191        model (tf.keras.Model): Trained model.192        tokenizer (Tokenizer): Fitted tokenizer.193        text_seq_len (int): Expected input sequence length.194        seed_text (str): Initial seed text.195        n_words (int): Number of words to generate.196 197    Returns:198        str: Generated text.199    """200    generated = []201    for _ in range(n_words):202        encoded = tokenizer.texts_to_sequences([seed_text])[0]203        encoded = pad_sequences([encoded], maxlen=text_seq_len, padding='pre')204        pred = model.predict(encoded, verbose=0)205        y_pred = np.argmax(pred, axis=1)[0]206        predicted_word = next((word for word, index in tokenizer.word_index.items() if index == y_pred), '')207        seed_text += ' ' + predicted_word208        generated.append(predicted_word)209    return ' '.join(generated)210 211def build_model(model_type: str, vocabulary_size: int, seq_length: int) -> tf.keras.Model:212    """213    Build and return a model based on the selected architecture.214 215    Args:216        model_type (str): Type of model architecture ('bi_di_gru', 'bi_di_lstm', 'gru', 'lstm').217        vocabulary_size (int): Size of the vocabulary.218        seq_length (int): Length of input sequences.219 220    Returns:221        tf.keras.Model: Compiled Keras model.222    """223    model = Sequential()224    model.add(Embedding(vocabulary_size, 50, input_length=seq_length))225    226    if model_type == "bi_di_gru":227        model.add(Bidirectional(GRU(100, return_sequences=True)))228        model.add(Dropout(0.2))229        model.add(GRU(120))230    elif model_type == "bi_di_lstm":231        model.add(Bidirectional(LSTM(100, return_sequences=True)))232        model.add(Dropout(0.2))233        model.add(LSTM(120))234    elif model_type == "gru":235        model.add(GRU(100, return_sequences=True))236        model.add(Dropout(0.2))237        model.add(GRU(120))238    elif model_type == "lstm":239        model.add(LSTM(100, return_sequences=True))240        model.add(Dropout(0.2))241        model.add(LSTM(120))242    else:243        raise ValueError("Invalid model type. Choose from 'bi_di_gru', 'bi_di_lstm', 'gru', or 'lstm'.")244    245    model.add(Dense(140, activation='relu'))246    model.add(Dense(vocabulary_size, activation='softmax'))247    248    model.build(input_shape=(None, seq_length))249    model.summary()250    251    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])252    return model253 254def main():255    # Set the model type here. Options: 'bi_di_gru', 'bi_di_lstm', 'gru', 'lstm'256    MODEL_TYPE = "bi_di_gru"257    258    # Prepare the data259    X, y, vocabulary_size, seq_length, tokenizer = prepare_data()260    261    # Build and train the selected model262    model = build_model(MODEL_TYPE, vocabulary_size, seq_length)263    history = model.fit(X, y, batch_size=512, epochs=300)264    265    # Save the model, tokenizer, and training history.266    model_filename = f"./{MODEL_TYPE.upper()}_model.h5"267    tokenizer_filename = f"{MODEL_TYPE}_tokenizer.pkl"268    history_filename = f"{MODEL_TYPE}_history.pkl"269    270    model.save(model_filename)271    print(f"Model saved to {model_filename}")272    273    with open(tokenizer_filename, 'wb') as handle:274        pickle.dump(tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL)275    print(f"Tokenizer saved to {tokenizer_filename}")276    277    with open(history_filename, 'wb') as file:278        pickle.dump(history.history, file)279    print(f"Training history saved to {history_filename}")280    281    # Plot training accuracy and loss282    epochs_range = range(len(history.history['accuracy']))283    plt.figure(figsize=(12, 5))284    plt.subplot(1, 2, 1)285    plt.plot(epochs_range, history.history['accuracy'])286    plt.title(f"{MODEL_TYPE.upper()} Model Training Accuracy")287    plt.xlabel("Epochs")288    plt.ylabel("Accuracy")289    290    plt.subplot(1, 2, 2)291    plt.plot(epochs_range, history.history['loss'])292    plt.title(f"{MODEL_TYPE.upper()} Model Training Loss")293    plt.xlabel("Epochs")294    plt.ylabel("Loss")295    plt.tight_layout()296    plt.show()297    298    # Generate stories from a few seed texts.299    seeds = [300        "The country was in chaos but",301        "I walked out of the store dissatisfied and it",302    ]303    304    for seed in seeds:305        generated_text = generate_story(model, tokenizer, seq_length, seed, 50)306        print("\nSeed:", seed)307        print("Generated Story:", generated_text)308 309if __name__ == "__main__":310    main()311 312