app.py

import streamlit as st
import streamlit.components.v1 as components
import pandas as pd
import os
import re
import numpy as np
from glob import glob
import lightgbm as lgb
import pickle

os.environ['S3_BUCKET'] = 'seriouslyusers'
os.environ['S3_BUCKET'] = "seriouslytestfaces"

import io

def get_s3_url(key):
    url = 'https://s3.amazonaws.com/%s/%s' % (os.environ['S3_BUCKET'],key.replace(' ','+'))
    return url

embeddings = pd.read_parquet('./embeddings.parquet')

def create_dir(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

def setup_user():
    create_dir(f'./users/{st.session_state.name}')
    create_dir(f'./users/{st.session_state.name}/likes')
    create_dir(f'./users/{st.session_state.name}/models')

def get_filename():
    if 'preds' in st.session_state:
        p = st.session_state.preds**4
        p /= sum(p)
        choice = np.random.choice(range(len(p)),p=p)
        st.session_state.pred = st.session_state.preds[choice]
        return embeddings.index[choice]
    st.toast('Random for now')
    return np.random.choice(embeddings.index)

st.title('What does attractive mean to you?')
st.session_state.name = st.text_input(label='Invent a unique alias (and remember it)')


def liked(filename,like):
    filename = f'./users/{st.session_state.name}/likes/' + filename.split('/')[-1] + '.' + str(like)[:1]
    open(filename, 'a').close()

def get_train_data():
    clean = lambda file : file.replace('\\','/').split('/')[-1][:-2]
    true_files = list(map(clean,glob(f'./users/{st.session_state.name}/likes/*.T')))
    false_files = list(map(clean,glob(f'./users/{st.session_state.name}/likes/*.F')))
    true_embeddings = embeddings.loc[true_files].values
    false_embeddings = embeddings.loc[false_files].values
    st.toast(f'Found {len(true_files)} positives and {len(false_files)} negatives')
    labels = np.array([1 for _ in true_embeddings] + [0 for _ in false_embeddings])
    st.session_state.labels = pd.Series(labels,index=true_files+false_files).rename('label')
    X = np.vstack([true_embeddings,false_embeddings])
    return X,labels

def train_model(X,labels):
    if len(labels) < 30:
        st.toast('Not enough data')
        return
    if labels.mean() > 0.9:
        st.toast('Not enough negatives')
        return
    if labels.mean() < 0.1:
        st.toast('Not enough positives')
        return
    train_data = lgb.Dataset(X, label=labels)
    num_round = 10
    param = {'num_leaves':10, 'objective': 'binary', 'metric' : 'auc'}
    bst = lgb.train(param, train_data, num_round)
    in_sample_preds = bst.predict(X)
    in_sample_score = np.corrcoef([in_sample_preds,np.array(labels)])[0][1]
    st.session_state.score = in_sample_score
    st.toast(f'Score = {in_sample_score:.1%}')
    return bst
    
def rank_candidates(bst):
    return bst.predict(embeddings.values)


def train():
    X,labels = get_train_data()
    bst = train_model(X,labels)
    if bst is None:
        return
    filename = f'./users/{st.session_state.name}/models/model.txt'
    bst.save_model(filename)
    preds = rank_candidates(bst)
    st.session_state.preds = preds


def cleanup():
    files = glob(f'./users/{st.session_state.name}/likes/*')
    for f in files:
       os.remove(f)
    if 'preds' in st.session_state:
        del st.session_state.preds
        del st.session_state.pred
        
    
def get_extremes(n=4):
    if 'preds' in st.session_state:
        preds = pd.Series(st.session_state.preds,index=embeddings.index).sort_values(ascending=False)
        return preds.iloc[:n].to_dict(),preds.iloc[-n:].to_dict()
        
def get_strange(n=4):
    if 'labels' in st.session_state:
        labels = st.session_state.labels
        preds = pd.Series(st.session_state.preds,index=embeddings.index).loc[labels.index].rename('pred')
        data = pd.concat([labels, preds],axis=1)
        st.toast(data.columns)
        
        data['diff'] = data['pred'] - data['label']
        data = data.sort_values('diff')['diff']
        return data.iloc[:n].to_dict(),data.iloc[-n:].to_dict() 
        
        


if st.session_state.name:
    st.session_state.name = re.sub(r'[^A-Za-z0-9 ]+', '', st.session_state.name)[:100]
    setup_user()
    st.subheader(f"Let's start {st.session_state.name}")
    filename = get_filename()
    
    
    cc1, cc2 = st.columns(2)
    c1,c2 = cc1.columns(2)
    c1.button('Why not', on_click=liked, args=[filename,True])
    c2.button('Nope', on_click=liked, args=[filename,False])
    key = get_s3_url(filename)
    cc1.image(key, width = 400)
    c1,c2 = cc2.columns(2)
    c1.button('Train',on_click=train,args=[])
    c2.button('Start over',on_click=cleanup,args=[])
    if 'preds' in st.session_state:
        
        cc1.write('Here is our guess')
        cc1.metric("Probability you will like", f'{st.session_state.pred:.1%}')
    
        best,worst = get_extremes()
        
        cc2.subheader('Predicted best')
        cs = cc2.columns(len(best))
        for c,(file,pred) in zip(cs,best.items()):
            c.metric("", f'{pred:.0%}')
            c.image(get_s3_url(file), width = 100)
            
        cc2.subheader('Predicted worst')
        cs = cc2.columns(len(worst))
        for c,(file,pred) in zip(cs,worst.items()):
            c.metric("", f'{pred:.0%}')
            c.image(get_s3_url(file), width = 100)

        cc1.metric("Overall model accuracy", f'{st.session_state.score:.0%}')

        cc1.subheader('Where you confused me')

        best,worst = get_strange()
        
        cc1.write("You didn't like my picks")
        cs = cc1.columns(len(best))
        for c,(file,pred) in zip(cs,best.items()):
            c.metric("", "",f'{pred:.0%}')
            c.image(get_s3_url(file), width = 100)

        cc1.write("You liked these more than I thought")
        cs = cc1.columns(len(worst))
        for c,(file,pred) in zip(cs,worst.items()):
            c.metric("","", f'{pred:.0%}')
            c.image(get_s3_url(file), width = 100)