<gh_stars>0
import streamlit as st
from transformers import pipeline
from datasets import load_dataset
import sys
import csv
import requests
import json

st.title('Test Streamlit')

st.header('Masked LMs')
@st.cache
def fill(model, text):
    unmasker = pipeline('fill-mask', model=model)
    return unmasker(text)

models = ['bert-base-uncased', 'bert-base-cased', 'roberta-base', 'distilbert-base-uncased']
model = st.selectbox('Select an MLM model', models)
text = st.text_input("Masked text", "i really like [MASK].")

if st.button('Fill'):
    st.table(fill(model, text))

st.header('Data Exploration')
dataset = load_dataset("sst")

left_col, right_col = st.beta_columns([1,1])

with left_col:
    label_counts = {True: {}, False: {}}
    all_insts = []
    search_tok = st.text_input("Enter left text", "nice")
    for typ,insts in dataset.items():
        label_counts[True][typ] = 0
        label_counts[False][typ] = 0
        for i in insts:
            if i['sentence'].find(search_tok) != -1:
                label = i['label']>0.5
                label_counts[label][typ] += 1
                all_insts.append({'review': i['sentence'], 'label':label})
    st.bar_chart(label_counts)
    # st.table(all_insts)
with right_col:
    label_counts = {True: {}, False: {}}
    all_insts = []
    search_tok = st.text_input("Enter right text", "nice")
    for typ,insts in dataset.items():
        label_counts[True][typ] = 0
        label_counts[False][typ] = 0
        for i in insts:
            if i['sentence'].find(search_tok) != -1:
                label = i['label']>0.5
                label_counts[label][typ] += 1
                all_insts.append({'review': i['sentence'], 'label':label})
    st.bar_chart(label_counts)
    # st.table(all_insts)    

st.header('GPT3')
def query(engine, text, api_token, num_completions):
    headers = {"Authorization": "Bearer " + api_token, "Content-Type": "application/json"}
    for i in range(num_completions):
        response = requests.post("https://api.openai.com/v1/engines/"+ engine +"/completions",
            headers=headers,
            json={
                'prompt': text,
                'max_tokens': 6,
            }
        ).json()['choices']
        st.markdown(text + " **" + response[0]['text']  + "**")

api_token = st.sidebar.text_input("OpenAI API token")
text = st.text_input("Enter a prompt", "I really like")
engine = st.selectbox('Select engine', ['ada', 'babbage', 'curie', 'davinci'])
num_completions = st.slider("Number of completions", 1, 10, 1)
query(engine, text, api_token, num_completions)