Refineverse_Plugin_backup / GenerationFeature.py
bsjghee001's picture
Duplicate from FYP-S1-21/Refineverse_Plugin
26f6649
import re # Python's built-in library for regular expressions (or Regex)
import sqlite3
from flask import g
from transformers import pipeline, set_seed
# Main function of the generation feature. Performs text generation!
def generate(Entered_story):
# Check if the input is empty
if not Entered_story.strip():
raise ValueError("Empty input!")
# Validate that the input is in the correct format
if not validate_story(Entered_story):
raise ValueError("Incorrect format!")
# Set the pipeline to use the correct NLP type and model
generator = pipeline('text-generation', model='gpt2')
# Take note: The max_length & min_length variables refer to the OUTPUT length!
set_seed(42)
generated_text = generator(Entered_story, max_length=30, num_return_sequences=5)
generated_text = generated_text[0]['generated_text']
return generated_text
# User Input Format Validation Function
def validate_story(Entered_story):
pattern = r'As a (?P<role>[^,.]+), I want to (?P<goal>[^,.]+)(?:,|.)+\s*so that' #Follows the normal structure, but allows anything after 'so that'
match = re.search(pattern, Entered_story, flags=re.DOTALL)
return bool(match)
# Function to grab all contents in the "TextGeneration" table (except for unique ids)
def getTextGenContents():
db = getattr(g, '_database', None) # Gets the _database attribute from the 'g' object. If it does not exist, returns 'None'
if db is None:
db = g._database = sqlite3.connect('Refineverse.db') # If db is None, create a new connection for db and g._database.
cursor = db.cursor() # Creates a cursor object to handle data
cursor.execute("SELECT userStory, generatedStory FROM TextGeneration") # The cursor executes the query
rows = cursor.fetchall() # Stores the results of fetchall() into a variable
return rows
# Function to insert a new row into the "TextGeneration" table
def insertTextGenRow( Entered_story, generatedStory):
with sqlite3.connect('Refineverse.db') as conn: # 'With' will automatically take care of closing and opening the connection
cursor = conn.cursor()
cursor.execute("INSERT INTO TextGeneration (userStory, generatedStory) VALUES (?, ?)", (Entered_story, generatedStory))
conn.commit()