# Contains only NLP zero-shot classification function codes! | |
import sqlite3 | |
from flask import g | |
from transformers import pipeline | |
# Set up zero-shot classification pipeline | |
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli") | |
# The main "breakdown" functionality. Performs NLP using zero-shot! | |
def breakdown(userStory): | |
# The results are stored into a dictionary variable with predifined labels | |
processedStory = classifier(userStory, candidate_labels=["developer", "tester", "project manager", "system admin"]) | |
# Extract labels and scores | |
scores = processedStory['scores'] | |
labels = processedStory['labels'] | |
# As the index of a score is always equal to its associated label, | |
# We only need to index of the score to find correct label. | |
maxScoreIndex = scores.index(max(scores)) # Gets the index of the highest score/accuracy | |
maxLabel = labels[maxScoreIndex] # Gets the associated label name | |
# Return labels | |
return maxLabel | |
# Function to grab all contents in the "Breakdown" table (except for unique ids) | |
# If adding any additional attributes to the table, this has to be updated accordingly | |
def getBreakdownContents(): | |
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 user_story, assignedLabel FROM Breakdown") # 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 "Breakdown" table | |
# Using "with" for the connection here seems important, as otherwise it results in an exception | |
def insertBreakdownRow(user_story, assigned_label): | |
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 Breakdown (user_story, assignedLabel) VALUES (?, ?)", (user_story, assigned_label)) | |
conn.commit() | |