Spaces:
Build error
Build error
| import streamlit as st | |
| import torch | |
| from pandas import options | |
| from transformers import BartForConditionalGeneration, BartTokenizer | |
| # initialize model + tok variables | |
| model = None | |
| tok = None | |
| # Examples for each models | |
| examples = [ | |
| ["interview-question-remake", "I have a cat named dolche and he's not very friendly with strangers. I've had him for 9 years now and it has been a pleasure to see him grow closer to us every year."], | |
| ["interview-length-tagged","Today's weather was really nice."], | |
| ["reverse-interview-question", "There are so many incredible musicians out there and so many really compelling big hits this year that it makes for a really interesting way to recap some of those big events."] | |
| ] | |
| # Descriptions for each models | |
| # descriptions = "Interview question remake is a model that..." | |
| # pass in Strings of model choice and input text for context | |
| def genQuestion(model_choice, context): | |
| # global descriptions | |
| if model_choice=="interview-question-remake": | |
| model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-question-remake") | |
| tok = BartTokenizer.from_pretrained("hyechanjun/interview-question-remake") | |
| # descriptions = "Interview question remake is a model that..." | |
| elif model_choice=="interview-length-tagged": | |
| model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-length-tagged") | |
| tok = BartTokenizer.from_pretrained("hyechanjun/interview-length-tagged") | |
| # descriptions = "Interview question tagged is a model that..." | |
| elif model_choice=="reverse-interview-question": | |
| model = BartForConditionalGeneration.from_pretrained("hyechanjun/reverse-interview-question") | |
| tok = BartTokenizer.from_pretrained("hyechanjun/reverse-interview-question") | |
| # descriptions = "Reverse interview question is a model that..." | |
| inputs = tok(context, return_tensors="pt") | |
| output = model.generate(inputs["input_ids"], num_beams=4, max_length=64, min_length=9, num_return_sequences=4, diversity_penalty =1.0, num_beam_groups=4) | |
| final_output = '' | |
| for i in range(4): | |
| final_output += [tok.decode(beam, skip_special_tokens=True, clean_up_tokenization_spaces=False) for beam in output][i] + "\n\n" | |
| return final_output | |
| # Wide page layout (instead of having a narrower, one-column page layout) | |
| st.set_page_config(layout="wide") | |
| # Title | |
| st.title("Interview AI Test Website") | |
| # Adding a Session State to store stateful variables and for saving user's labels/tags for generated questions | |
| if 'button_sent' not in st.session_state: | |
| st.session_state.button_sent = False | |
| # Input fields | |
| input = st.text_input('Context') # user inputs context to construct a response (str) | |
| maxl, minl = st.columns(2) | |
| option = st.selectbox( | |
| 'Please select a model.', | |
| ('interview-question-remake', 'interview-length-tagged', 'reverse-interview-question')) | |
| if option == 'interview-question-remake': | |
| st.write("This is the re-fine-tuned base model for our interview AI. It returns strings terminating in a question mark (?).") | |
| elif option == 'interview-length-tagged': | |
| st.write("This is a length-tagged version of our interview AI. You can specify how long its responses should be (ranges of multiples of 10)") | |
| elif option == 'reverse-interview-question': | |
| st.write("This model asks a question that would have resulted in the context you provide (a.k.a. it traverses backward through the interview)") | |
| # Column layout to display generated responses alongside tags | |
| col1, col2 = st.columns((3, 1)) | |
| if st.button('Submit') or st.session_state.button_sent: | |
| with st.spinner('Generating a response...'): | |
| output = genQuestion(option, input) | |
| print(output) | |
| # st.write(output) | |
| st.session_state.button_sent = True | |
| col1.text_area(label="Generated Responses:", value=output, height=200) | |
| # TODO: | |
| # - disable multiselect widget when responses are being generated AND when a question is not selected to be tagged | |
| # - connect tags with an individual question | |
| # - save session state so tags associated with their respective questions can also be saved | |
| # - write/store the saved state data to some database for future use? | |
| # - brainstorm good names for tags/labels OR allow users to enter their own tag names if possible | |