Refineverse_Plugin_testing / BreakdownFeature.py
bsjghee001's picture
Duplicate from FYP-S1-21/Refineverse_Plugin
e7942cf
# 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()