app.py

# Imports
# Core Imports
import torch

# Model-related Imports
from transformers import BartTokenizer, BartForConditionalGeneration # fine-tuned BART model
from transformers import AutoTokenizer, AutoModelForTokenClassification # restore punct
from transformers import pipeline # restore punct
import gradio as gr

# Evaluation Imports
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
import string



# Instantiate model to restore punctuation
print("1/7 - Instantiating model to restore punctuation")

punct_model_path = "felflare/bert-restore-punctuation"
# Load punct tokenizer and model
punct_tokenizer = AutoTokenizer.from_pretrained(punct_model_path)
punct_model = AutoModelForTokenClassification.from_pretrained(punct_model_path)
punct_restorer = pipeline("token-classification", model=punct_model, tokenizer=punct_tokenizer)



# Instantiate fine-tuned horror BART model
print("2/7 - Instantiating two-sentence horror generation model")

model_path = 'voacado/bart-two-sentence-horror'
# Load tokenizer and model
tokenizer = BartTokenizer.from_pretrained(model_path)
model = BartForConditionalGeneration.from_pretrained(model_path)



# Load data for evaluation metrics
print("3/7 - Reading in data")
data = pd.read_csv("./reddit_cleansed_data.csv")
data['weighted_score'] = data['score'] + (10 * data['num_comments']) + (100 * data['gilded_count']) 
dataset_stories = (data['title'] + ' ' + data['selftext']).to_list()



# Instantiate evaluation metrics - Cosine Similarity with TF-IDF
print("4/7 - Instantiating evaluation metrics - Cosine Similarity with TF-IDF")
# Pre-vectorize dataset
vectorizer = TfidfVectorizer()
dataset_matrix = vectorizer.fit_transform(dataset_stories)

def eval_cosine_similarity(input_sentence: str) -> [str, str]:
    """
    Evaluate cosine similarity between input sentence and each story in the dataset.

    Args:
        input_sentence (str): user story (first sentence)

    Returns:
        [str, str]: most similar story, weighted score
    """
    # Vectorize input sentence using the existing vocab
    input_vec = vectorizer.transform([input_sentence])
    # Get cosine similarity
    similarities = cosine_similarity(input_vec, dataset_matrix)
    # Find most similar story
    most_similar_story_idx = similarities.argmax()
    most_similar_story = dataset_stories[most_similar_story_idx]
    # Get weighted score of most similar story
    weighted_score = data['weighted_score'][most_similar_story_idx]

    return most_similar_story, weighted_score



# Instantiate evaluation metrics - Jaccard Similarity
print("5/7 - Instantiating evaluation metrics - Jaccard Similarity")
def tokenize(text: str):
    """
    Convert text to lowercase and remove punctuation, then tokenize.

    Args:
        text (str): user story

    Returns:
        set: set of tokens
    """
    text = text.lower()
    text = text.translate(str.maketrans('', '', string.punctuation))
    tokens = text.split()
    return set(tokens)

def jaccard_similarity(set1: set, set2: set):
    """
    Calculate Jaccard similarity between two sets.

    Args:
        set1 (set): user_tokens
        set2 (set): story_tokens

    Returns:
        float: Jaccard similarity
    """
    intersection = set1.intersection(set2)
    union = set1.union(set2)
    return len(intersection) / len(union)

def eval_jaccard_similarity(input_sentence: str) -> [str, str]:
    """
    Evaluate Jaccard similarity between input sentence and each story in the dataset.

    Args:
        input_sentence (str): user story (first sentence)

    Returns:
        [str, str]: most similar story, weighted score
    """
    # Tokenize the user story
    user_tokens = tokenize(input_sentence)

    # Initialize variables to find the most similar story
    max_similarity = 0
    most_similar_story = ''

    # Compare with each story in the dataset
    for story in dataset_stories:
        story_tokens = tokenize(story)
        similarity = jaccard_similarity(user_tokens, story_tokens)
        if similarity > max_similarity:
            max_similarity = similarity
            most_similar_story = story
            max_score = data['weighted_score'][dataset_stories.index(story)]
    
    return most_similar_story, max_score



# Set up inference
print("6/7 - Setting parameters for inference")

# Set the model to evaluation mode
model.eval()
# If GPU, use it
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Restore punct
def restore_punctuation(text: str, restorer: pipeline) -> str:
    """
    Restore punctuation to text.

    Args:
        text (str): full story (first and second sentences)
        restorer (pipeline): model that restores punctuation

    Returns:
        str: punctuated text (based on input)
    """
    # Use the model to predict punctuation
    punctuated_output = restorer(text)
    punct_text = []
    
    # Define punctuation marks (note: not including left-side because we want space still)
    punctuation_marks = ["!", "?", ".", "-", ":", ";", "'", "’", ",", ")", "]", "}", "…", "”", "’’", "''"]
    
    for elem in punctuated_output:
        cur_token = elem.get('word')
        
        # If token is punctuation, append to previous token
        if cur_token in punctuation_marks:
            punct_text[-1] += cur_token
            
        # If previous token is quotations, append to previous token
        elif punct_text and punct_text[-1] in ["'", "’", "“", "‘", "‘‘", "““"]:
            punct_text[-1] += cur_token
            
        # If token is a contraction or a quote, append to previous token (no space)
        elif cur_token.lower() in ["s", "t", "re", "ve", "ll", "d", "m"]:
            # Remove space for contractions
            punct_text[-1] += cur_token
            
        # if prediction is LABEL_0, token should be capitalized
        elif elem.get('entity') == 'LABEL_0':
            punct_text.append(cur_token.capitalize())

        # else if prediction is LABEL_1, token should be lowercase
        # elif elem.get('entity') == 'LABEL_1':
        else:
            punct_text.append(cur_token)
            
    # If there's no period at the end of the story, add one
    if punct_text[-1][-1] != '.':
        punct_text[-1] = punct_text[-1] + '.'

    return ' '.join(punct_text)

def generate_text(input_text: str, full_sentence: str) -> [str, str, float, str, float]:
    """
    Generate the second sentence of the horror story given the first (input_text).

    Args:
        input_text (str): first sentence of the horror story
        full_sentence (str): full story (first and second sentences)

    Returns:
        gen_text_punct (str): second sentence of the horror story
        similar_story_cosine (str): most similar story (cosine similarity)
        cosine_score (float): score of most similar story (cosine similarity)
        similar_story_jaccard (str): most similar story (Jaccard similarity)
        jaccard_score (float): score of most similar story (Jaccard similarity)
    """
    # If user only enters first sentence, generate second sentence
    if not full_sentence:
        # Encode the input text
        input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)

        # Generate text
        with torch.no_grad():
            output_ids = model.generate(input_ids, max_length=50)

        # Decode the generated text
        gen_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        
        # Restore punctuation
        gen_text_punct = restore_punctuation(gen_text, punct_restorer)
        full_sentence = input_text + ' ' + gen_text_punct
    else:
        gen_text_punct = "N/A"
    
    # Calculate Cosine and Jaccard similarity
    similar_story_cosine, cosine_score = eval_cosine_similarity(full_sentence)
    similar_story_jaccard, jaccard_score = eval_jaccard_similarity(full_sentence)

    return gen_text_punct, similar_story_cosine, cosine_score, similar_story_jaccard, jaccard_score



# Create gradio demo
print("7/7 - Launching demo")

title = "👻 🫣 Generate a Two-Sentence Horror Story 😱 👻"
description = """
<center>The bot was trained to generate two-sentence horror stories based on r/TwoSentenceHorror. <i>Spooky!</i></center>
"""

article = """
Check out [the subreddit](https://www.reddit.com/r/TwoSentenceHorror) that this demo is based off of. Or, check out the dataset [here](https://www.kaggle.com/datasets/voanthony/two-sentence-horror-jan-2015-apr-2023).

The language model is fine-tuned from ['facebook/bart-base'](https://huggingface.co/facebook/bart-base). We import, then update the weights for the model to generate two-sentence horror stories. The model is fine-tuned over 3 epochs to avoid catastrophic forgetting. We also use a separate model (['felflare/bert-restore-punctuation'](https://huggingface.co/felflare/bert-restore-punctuation?text=My+name+is+wolfgang+and+I+live+in+berlin)) to restore punctuation.

For evaluation, the generated story is compared to the most similar Reddit post (using either cosine or Jaccard similarity). The score of the most similar post is also returned. The score is calculated as the sum of the post score, 10 * number of comments, and 100 * number of gilds. The score is used as a proxy for the popularity of the post.

Users may also enter an entire story in the second input prompt rather than generating the remainder of the story. This will be used for evaluation metrics and no story will be generated.
"""


demo = gr.Interface(
    fn=generate_text, 
    inputs=[
        gr.Textbox(lines=4, placeholder="Enter the first sentence of your horror story here...", label="First Sentence"),
        gr.Textbox(lines=4, placeholder="Or, enter full story for evaluation here...", label="Eval - Full Story")
    ],
    outputs=[
        gr.Textbox(lines=4, label="Generated Second Sentence"),
        gr.Textbox(lines=3, label="Cosine Similarity - Sentence"),
        gr.Textbox(lines=1, label="Cosine Similarity - Post Score"),
        gr.Textbox(lines=3, label="Jaccard Similarity - Sentence"),
        gr.Textbox(lines=1, label="Jaccard Similarity - Post Score")
    ],
    title=title,
    description=description,
    article=article,
    examples=[["My parents told me not to go upstairs."], ["There was a ghost."], ["Sometimes I catch myself staring at those missing person flyers at the store."]],
    )

demo.launch(share=True)