add connections to database

Home.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 import streamlit as st
 import random
 import time
@@ -28,7 +29,7 @@ def save_user_id(user_id):
     print(user_id)
     if not user_id:
         user_id = 'anonymous' + str(random.randint(0, 100000))
-    st.session_state.user_id = [user_id, time.time()]
+    st.session_state.user_id = [user_id, datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
 
 
 def logout():

data/ranking_script.py

@@ -1,4 +1,5 @@
 from datasets import Dataset
+from datetime import datetime
 
 
 def init_ranking_data():
@@ -6,7 +7,7 @@ def init_ranking_data():
 
     # add example data
     # note that image_id is a string, other ids are int
-    ds = ds.add_item({'image_id': '0', 'modelVersion_id': 0, 'ranking': 0, "user_name": "example_data", "timestamp": 0.0})
+    ds = ds.add_item({'image_id': '0', 'modelVersion_id': 0, 'ranking': 0, "user_name": "example_data", "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")})
 
     ds.push_to_hub("MAPS-research/GEMRec-Ranking", split='train')
 

pages/Ranking.py

@@ -1,6 +1,9 @@
+import os
+
 import datasets
 import numpy as np
 import pandas as pd
+import pymysql.cursors
 import streamlit as st
 
 from streamlit_elements import elements, mui, html, dashboard, nivo
@@ -165,22 +168,28 @@ class RankingApp:
 
             restart_btn = st.button('🎖️ Rank Again')
             if restart_btn:
-                st.session_state.progress['prompt_id'] = 'ranking'
+                st.session_state.progress[prompt_id] = 'ranking'
                 st.session_state.counter[prompt_id] = 0
                 st.experimental_rerun()
 
 
     def next_batch(self, prompt_id, progress=None):
-
         # save ranking to dataset
         # print(st.session_state.ranking)
-        ranking_dataset = datasets.load_dataset('MAPS-research/GEMRec-Ranking', split='train')
+        # ranking_dataset = datasets.load_dataset('MAPS-research/GEMRec-Ranking', split='train')
+        curser = RANKING_CONN.cursor()
         for image_id in st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]].keys():
             modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_id]['modelVersion_id'].values[0]
             ranking = st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][image_id]
             # print({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
-            ranking_dataset = ranking_dataset.add_item({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
-        ranking_dataset.push_to_hub('MAPS-research/GEMRec-Ranking', split='train')
+            # ranking_dataset = ranking_dataset.add_item({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
+
+            query = "INSERT INTO rankings (image_id, modelVersion_id, ranking, user_name, timestamp) VALUES (%s, %s, %s, %s, %s)"
+            curser.execute(query, (image_id, modelVersion_id, ranking, st.session_state.user_id[0], st.session_state.user_id[1]))
+
+        curser.close()
+        RANKING_CONN.commit()
+        # ranking_dataset.push_to_hub('MAPS-research/GEMRec-Ranking', split='train')
 
         if progress == 'finished':
             st.session_state.progress[prompt_id] = 'finished'
@@ -191,6 +200,20 @@ class RankingApp:
         st.session_state.counter[prompt_id] -= 1
 
 
+def connect_to_db():
+    conn = pymysql.connect(
+        host=os.environ.get('RANKING_DB_HOST'),
+        port=3306,
+        database='myRanking',
+        user=os.environ.get('RANKING_DB_USER'),
+        password=os.environ.get('RANKING_DB_PASSWORD'),
+        charset='utf8mb4',
+        cursorclass=pymysql.cursors.DictCursor
+    )
+
+    return conn
+
+
 if __name__ == "__main__":
     st.set_page_config(page_title="Personal Image Ranking", page_icon="🎖️️", layout="wide")
 
@@ -218,7 +241,9 @@ if __name__ == "__main__":
             # st.write('You have checked ' + str(len(selected_modelVersions)) + ' images.')
             roster, promptBook, images_ds = load_hf_dataset()
             print(st.session_state.selected_dict)
+
             # st.write("# Full function is coming soon.")
+            RANKING_CONN = connect_to_db()
 
             # only select the part of the promptbook where tag is the same as st.session_state.selected_dict.keys(), while model version ids are the same as corresponding values to each key
             promptBook_selected = pd.DataFrame()

requirements.txt

@@ -3,4 +3,5 @@ streamlit-elements==0.1.0
 streamlit-extras
 altair<5
 streamlit-vega-lite
-scikit-learn
+scikit-learn
+pymysql