Update app.py

app.py

@@ -10,7 +10,36 @@ from elo import update_elo_ratings  # Custom function for ELO ratings
 enable_btn = gr.Button.update(interactive=True)
 
 import sqlite3
+import requests
 
+def classify_vote(user_input):
+    url = "https://api-inference.huggingface.co/models/facebook/bart-large-mnli"
+    headers = {
+        "accept": "*/*",
+        "accept-language": "en-US,en;q=0.9",
+        "content-type": "application/json",
+        "sec-ch-ua": "\"Not A(Brand\";v=\"99\", \"Microsoft Edge\";v=\"121\", \"Chromium\";v=\"121\"",
+        "sec-ch-ua-mobile": "?0",
+        "sec-ch-ua-platform": "\"macOS\"",
+        "sec-fetch-dest": "empty",
+        "sec-fetch-mode": "cors",
+        "sec-fetch-site": "same-site"
+    }
+    payload = {
+        "inputs": user_input,
+        "parameters": {
+            "candidate_labels": "information and knowledge,problem solving,creative expression,personal development and guidance,entertainment, niche topics",
+            "multi_class": True
+        }
+    }
+    response = requests.post(url, headers=headers, json=payload)
+    if response.status_code == 200:
+        response_data = response.json()
+        top_category = response_data["labels"][0].strip()
+        return top_category
+    else:
+        print(f"Error: {response.status_code}")
+        return None
 from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
 async def direct_regenerate(model, user_input, chatbot, character_name, character_description, user_name):
@@ -28,21 +57,10 @@ password=os.environ.get("MONGODB")
 def init_database():
     uri = f"mongodb+srv://new-user:{password}@cluster0.xb2urf6.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0"
     client = MongoClient(uri)
-    db = client["elo_ratings"]
-    collection = db["elo_ratings"]
+    db = client["elo_ratings2"]
+    collection = db["elo_ratings2"]
     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)
-
 import json
 
 with open('chatbots.txt', 'r') as file:
@@ -142,7 +160,7 @@ async def chat_with_bots(user_input, state, character_name, character_descriptio
         get_bot_response(bot2_adapter, user_input, state, 1, character_name, character_description, user_name)
     )
     return bot1_response.replace("<|im_end|>",""), bot2_response.replace("<|im_end|>","")
-def update_ratings(state, winner_index, collection):
+def update_ratings(state, winner_index, collection, category):
     elo_ratings = get_user_elo_ratings(collection)
     winner_adapter = state['last_bots'][winner_index]
     loser_adapter = state['last_bots'][1 - winner_index]
@@ -150,24 +168,28 @@ def update_ratings(state, winner_index, collection):
     winner = next(entry['original_model'] for entry in chatbots_data if entry['adapter'] == winner_adapter)
     loser = next(entry['original_model'] for entry in chatbots_data if entry['adapter'] == loser_adapter)
     
-    elo_ratings = update_elo_ratings(elo_ratings, winner_adapter, loser_adapter)
-    update_elo_rating(collection, elo_ratings, winner_adapter, loser_adapter)
+    elo_ratings = update_elo_ratings(elo_ratings, winner_adapter, loser_adapter, category)
+    update_elo_rating(collection, elo_ratings, winner_adapter, loser_adapter, category)
     return [('Winner: ', winner), ('Loser: ', loser)]
     
-def vote_up_model(state, chatbot, chatbot2):
+def vote_up_model(state, chatbot, chatbot2, user_input):
     collection = init_database()
-    update_message = update_ratings(state, 0, collection)
-    chatbot.append(update_message[0])
-    chatbot2.append(update_message[1])
+    top_category = classify_vote(user_input)
+    if top_category:
+        update_message = update_ratings(state, 0, collection, top_category)
+        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)
 
-def vote_down_model(state, chatbot, chatbot2):
+def vote_down_model(state, chatbot, chatbot2, user_input):
     collection = init_database()
-    update_message = update_ratings(state, 1, collection)
-    chatbot2.append(update_message[0])
-    chatbot.append(update_message[1])
+    top_category = classify_vote(user_input)
+    if top_category:
+        update_message = update_ratings(state, 1, collection, top_category)
+        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)
-
+    
 async def user_ask(state, chatbot1, chatbot2, textbox, character_name, character_description, user_name):
     if character_name and len(character_name) > 20:
         character_name = character_name[:20]  # Limit character name to 20 characters
@@ -230,16 +252,39 @@ def submit_model(model_name):
     else:
         return "Discord webhook URL not configured."
         
+def get_user_elo_ratings(collection):
+    rows = list(collection.find())
+    if rows:
+        elo_ratings = {}
+        for row in rows:
+            bot_name = row['bot_name']
+            if bot_name not in elo_ratings:
+                elo_ratings[bot_name] = {}
+            for category in row['categories']:
+                elo_ratings[bot_name][category] = {'elo_rating': row['categories'][category]['elo_rating'], 'games_played': row['categories'][category]['games_played']}
+        return elo_ratings
+    else:
+        return {"default": {'overall': {'elo_rating': 1200, 'games_played': 0}}}
+
+def update_elo_rating(collection, updated_ratings, winner, loser, category):
+    collection.update_one({"bot_name": winner}, {"$set": {f"categories.{category}.elo_rating": updated_ratings[winner][category]['elo_rating'], f"categories.{category}.games_played": updated_ratings[winner][category]['games_played']}}, upsert=True)
+    collection.update_one({"bot_name": loser}, {"$set": {f"categories.{category}.elo_rating": updated_ratings[loser][category]['elo_rating'], f"categories.{category}.games_played": updated_ratings[loser][category]['games_played']}}, upsert=True)
+
 def generate_leaderboard(collection):
     rows = list(collection.find())
-    leaderboard_data = pd.DataFrame(rows, columns=['bot_name', 'elo_rating', 'games_played'])
-    leaderboard_data['original_model'] = leaderboard_data['bot_name'].apply(lambda x: next(entry['original_model'] for entry in chatbots_data if entry['adapter'] == x))
-    leaderboard_data = leaderboard_data[['original_model', 'elo_rating', 'games_played']]
-    leaderboard_data.columns = ['Chatbot', 'ELO Score', 'Games Played']
-    leaderboard_data['ELO Score'] = leaderboard_data['ELO Score'].round().astype(int)
-    leaderboard_data = leaderboard_data.sort_values('ELO Score', ascending=False)
-    return leaderboard_data
+    leaderboard_data = []
+    for row in rows:
+        bot_name = row['bot_name']
+        original_model = next(entry['original_model'] for entry in chatbots_data if entry['adapter'] == bot_name)
+        total_elo = sum(category['elo_rating'] for category in row['categories'].values())
+        total_games = sum(category['games_played'] for category in row['categories'].values())
+        avg_elo = total_elo / len(row['categories']) if len(row['categories']) > 0 else 0
+        leaderboard_data.append([original_model, avg_elo, total_games])
     
+    leaderboard_data = pd.DataFrame(leaderboard_data, columns=['Chatbot', 'Avg ELO Score', 'Total Games Played'])
+    leaderboard_data['Avg ELO Score'] = leaderboard_data['Avg ELO Score'].round().astype(int)
+    leaderboard_data = leaderboard_data.sort_values('Avg ELO Score', ascending=False)
+    return leaderboard_data
 def refresh_leaderboard():
     collection = init_database()
     leaderboard_data = generate_leaderboard(collection)
@@ -308,9 +353,8 @@ with gr.Blocks() as demo:
         textbox.submit(user_ask, inputs=[state, chatbot1, chatbot2, textbox, character_name, character_description, user_name], outputs=[state, chatbot1, chatbot2, textbox, upvote_btn_a, upvote_btn_b], queue=True)
         collection = init_database()
 
-        upvote_btn_a.click(vote_up_model, inputs=[state, chatbot1, chatbot2], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
-        upvote_btn_b.click(vote_down_model, inputs=[state, chatbot1, chatbot2], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
-
+        upvote_btn_a.click(vote_up_model, inputs=[state, chatbot1, chatbot2, textbox], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
+        upvote_btn_b.click(vote_down_model, inputs=[state, chatbot1, chatbot2, textbox], outputs=[chatbot1, chatbot2, upvote_btn_a, upvote_btn_b, textbox, submit_btn])
     with gr.Tab("💬 Direct Chat"):
         gr.Markdown("## 🗣️ Chat directly with a model!")