from nltk import download
from nltk.tokenize import sent_tokenize
from setfit import SetFitModel
import gradio as gr
import os

download('punkt')

inf_model = SetFitModel.from_pretrained("MYox/NSS_multi-label_8_sent_v3_1")

features = (['Teaching & learning', 'Support', 'Communication', 'Organisation & timetable', 
    'Assessment & feedback', 'Career & placement', 'Health, wellbeing & social life', 'Facilities & technology'])

def get_doc_topics(d):
    t = inf_model(d)
    return [f for i, f in enumerate(features) if t[0][i]]

def get_comment_topics(comment):
    sent_topics = []
    for s in sent_tokenize(comment):
        sent_topics.append(get_doc_topics([s]))

    output = []
    for x in sent_topics:
        for y in x:
            if y not in output:
                output.append(y)
    return '; '.join(output) 

app = gr.Interface(
    title='NSS Topic Generator',
    fn=get_comment_topics,
    inputs=gr.Textbox(lines=4, placeholder='Type comment here or choose example, below', label='Comment'),
    outputs=gr.Textbox(lines=4, label='Generated Topics'),
    allow_flagging='never',
    examples = ([["Lots of support, especially when revising for exams. Lecturers always respond to emails very quickly!"], 
        ["Loved everything. Met lots of new people and really enjoyed placements."],
        ["I thought the lectures were sometimes boring but the library was great!"]
        ])
)

app.launch()