Update app.py

app.py

@@ -9,15 +9,24 @@ enable_btn = gr.Button.update(interactive=True)
 
 import sqlite3
 
+import pymongo
+
 def init_database():
-    conn = sqlite3.connect('elo_ratings.db')
-    c = conn.cursor()
-    c.execute('''CREATE TABLE IF NOT EXISTS elo_ratings
-                 (bot_name TEXT PRIMARY KEY,
-                  elo_rating INTEGER,
-                  games_played INTEGER)''')
-    conn.commit()
-    conn.close()
+    client = pymongo.MongoClient("mongodb+srv://rwitzman:<password>@cluster0.xb2urf6.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0")
+    db = client["elo_ratings"]
+    collection = db["elo_ratings"]
+    return collection
+
+def get_user_elo_ratings(collection):
+    rows = list(collection.find())
+    if rows:
+        return {row['bot_name']: {'elo_rating': row['elo_rating'], 'games_played': row['games_played']} for row in rows}
+    else:
+        return {"default": {'elo_rating': 1200, 'games_played': 0}}
+
+def update_elo_rating(collection, updated_ratings, winner, loser):
+    collection.update_one({"bot_name": winner}, {"$set": {"elo_rating": updated_ratings[winner]['elo_rating'], "games_played": updated_ratings[winner]['games_played']}}, upsert=True)
+    collection.update_one({"bot_name": loser}, {"$set": {"elo_rating": updated_ratings[loser]['elo_rating'], "games_played": updated_ratings[loser]['games_played']}}, upsert=True)
 
 # Load chatbot URLs and model names from a JSON file
 # Load chatbot model adapter names from a text file
@@ -39,25 +48,7 @@ from datasets import load_dataset,DatasetDict,Dataset
 import requests
 import os
 
-def get_user_elo_ratings():
-    conn = sqlite3.connect('elo_ratings.db')
-    c = conn.cursor()
-    c.execute("SELECT * FROM elo_ratings")
-    rows = c.fetchall()
-    conn.close()
-    if rows:
-        return {row[0]: {'elo_rating': row[1], 'games_played': row[2]} for row in rows}
-    else:
-        return {"default": {'elo_rating': 1200, 'games_played': 0}}
-def update_elo_rating(updated_ratings, winner, loser):
-    conn = sqlite3.connect('elo_ratings.db')
-    c = conn.cursor()
-    c.execute("INSERT OR REPLACE INTO elo_ratings (bot_name, elo_rating, games_played) VALUES (?, ?, ?)",
-              (winner, updated_ratings[winner]['elo_rating'], updated_ratings[winner]['games_played']))
-    c.execute("INSERT OR REPLACE INTO elo_ratings (bot_name, elo_rating, games_played) VALUES (?, ?, ?)",
-              (loser, updated_ratings[loser]['elo_rating'], updated_ratings[loser]['games_played']))
-    conn.commit()
-    conn.close()
+
 # Function to get bot response
 def format_alpaca_prompt(state):
     alpaca_prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
@@ -124,21 +115,22 @@ async def chat_with_bots(user_input, state):
     )
 
     return bot1_response, bot2_response
-def update_ratings(state, winner_index):
-    elo_ratings = get_user_elo_ratings()
+def update_ratings(state, winner_index, collection):
+    elo_ratings = get_user_elo_ratings(collection)
     winner = state['last_bots'][winner_index]
     loser = state['last_bots'][1 - winner_index]
     
     elo_ratings = update_elo_ratings(elo_ratings, winner, loser)
-    update_elo_rating(elo_ratings, winner, loser)
-    return [('Winner: ', winner.replace('rwitz/','').replace('-lora','')), ('Loser: ', lose.replace('rwitz/','').replace('-lora',''))]
-def vote_up_model(state, chatbot,chatbot2):
-    update_message = update_ratings(state, 0)
+    update_elo_rating(collection, elo_ratings, winner, loser)
+    return [('Winner: ', winner.replace('rwitz/','').replace('-lora','')), ('Loser: ', loser.replace('rwitz/','').replace('-lora',''))]
+    
+def vote_up_model(state, chatbot, chatbot2, collection):
+    update_message = update_ratings(state, 0, collection)
     chatbot.append(update_message[0])
     chatbot2.append(update_message[1])
     return chatbot, chatbot2,gr.Button.update(interactive=False),gr.Button.update(interactive=False),gr.Textbox.update(interactive=False),gr.Button.update(interactive=False)  # Disable voting buttons
-def vote_down_model(state, chatbot,chatbot2):
-    update_message = update_ratings(state, 1)
+def vote_down_model(state, chatbot, chatbot2, collection):
+    update_message = update_ratings(state, 1, collection)
     chatbot2.append(update_message[0])
     chatbot.append(update_message[1])
     return chatbot,chatbot2, gr.Button.update(interactive=False),gr.Button.update(interactive=False),gr.Textbox.update(interactive=False),gr.Button.update(interactive=False)  # Disable voting buttons
@@ -179,13 +171,10 @@ import pandas as pd
 
 # Function to generate leaderboard data
 
-def generate_leaderboard():
-    conn = sqlite3.connect('elo_ratings.db')
-    c = conn.cursor()
-    c.execute("SELECT bot_name, elo_rating, games_played FROM elo_ratings ORDER BY elo_rating DESC")
-    rows = c.fetchall()
-    conn.close()
-    leaderboard_data = pd.DataFrame(rows, columns=['Chatbot', 'ELO Score', 'Games Played'])
+def generate_leaderboard(collection):
+    rows = list(collection.find())
+    leaderboard_data = pd.DataFrame(rows, columns=['bot_name', 'elo_rating', 'games_played'])
+    leaderboard_data.columns = ['Chatbot', 'ELO Score', 'Games Played']
     return leaderboard_data
     
 def refresh_leaderboard():
@@ -227,5 +216,7 @@ with gr.Blocks() as demo:
 
     # Launch the Gradio interface
 if __name__ == "__main__":
-    init_database()
+    collection = init_database()
+    upvote_btn_a.click(vote_up_model, inputs=[state, chatbot1, chatbot2, collection], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
+    upvote_btn_b.click(vote_down_model, inputs=[state, chatbot1, chatbot2, collection], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
     demo.launch(share=False)