Update app.py

app.py

@@ -27,6 +27,10 @@ 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')
+    if 'count' not in st.session_state:
+        st.session_state.count = 0
+        st.session_state.neg = 0
+        st.session_state.pos = 0
 
 def get_filename():
     if 'preds' in st.session_state:
@@ -45,6 +49,11 @@ st.session_state.name = st.text_input(label='Invent a unique alias (and remember
 def liked(filename,like):
     filename = f'./users/{st.session_state.name}/likes/' + filename.split('/')[-1] + '.' + str(like)[:1]
     open(filename, 'a').close()
+    st.session_state.count += 1
+    if like:
+        st.session_state.pos += 1
+    else:
+        st.session_state.neg += 1
 
 def get_train_data():
     clean = lambda file : file.replace('\\','/').split('/')[-1][:-2]
@@ -93,13 +102,17 @@ def train():
     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
+    for var in 'preds pred count pos neg'.split():        
+        if var in st.session_state:
+            del st.session_state[var]
+    st.session_state.count = 0
+    st.session_state.neg = 0
+    st.session_state.pos = 0
         
     
 def get_extremes(n=4):
@@ -126,6 +139,7 @@ def get_strange(n=4):
 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()
     
@@ -137,7 +151,8 @@ if st.session_state.name:
     key = get_s3_url(filename)
     cc1.image(key, width = 400)
     c1,c2 = cc2.columns(2)
-    c1.button('Train',on_click=train,args=[])
+    if st.session_state.count>30 and st.session_state.pos > 5 and st.session_state.neg > 5:
+        c1.button('Train',on_click=train,args=[])
     c2.button('Start over',on_click=cleanup,args=[])
     if 'preds' in st.session_state: