Spaces:
Running
Running
| import numpy as np | |
| import json | |
| from trueskill import TrueSkill | |
| import paramiko | |
| import io, os | |
| import sys | |
| import random | |
| sys.path.append('../') | |
| from serve.constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_SKILL | |
| trueskill_env = TrueSkill() | |
| ssh_matchmaker_client = None | |
| sftp_matchmaker_client = None | |
| def create_ssh_matchmaker_client(server, port, user, password): | |
| global ssh_matchmaker_client, sftp_matchmaker_client | |
| ssh_matchmaker_client = paramiko.SSHClient() | |
| ssh_matchmaker_client.load_system_host_keys() | |
| ssh_matchmaker_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) | |
| ssh_matchmaker_client.connect(server, port, user, password) | |
| transport = ssh_matchmaker_client.get_transport() | |
| transport.set_keepalive(60) | |
| sftp_matchmaker_client = ssh_matchmaker_client.open_sftp() | |
| def is_connected(): | |
| global ssh_matchmaker_client, sftp_matchmaker_client | |
| if ssh_matchmaker_client is None or sftp_matchmaker_client is None: | |
| return False | |
| if not ssh_matchmaker_client.get_transport().is_active(): | |
| return False | |
| try: | |
| sftp_matchmaker_client.listdir('.') | |
| except Exception as e: | |
| print(f"Error checking SFTP connection: {e}") | |
| return False | |
| return True | |
| def ucb_score(trueskill_diff, t, n): | |
| exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5)) | |
| ucb = -trueskill_diff + 1.0 * exploration_term | |
| return ucb | |
| def update_trueskill(ratings, ranks): | |
| new_ratings = trueskill_env.rate(ratings, ranks) | |
| return new_ratings | |
| def serialize_rating(rating): | |
| return {'mu': rating.mu, 'sigma': rating.sigma} | |
| def deserialize_rating(rating_dict): | |
| return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma']) | |
| def save_json_via_sftp(ratings, comparison_counts, total_comparisons): | |
| global sftp_matchmaker_client | |
| if not is_connected(): | |
| create_ssh_matchmaker_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| data = { | |
| 'ratings': [serialize_rating(r) for r in ratings], | |
| 'comparison_counts': comparison_counts.tolist(), | |
| 'total_comparisons': total_comparisons | |
| } | |
| json_data = json.dumps(data) | |
| with sftp_matchmaker_client.open(SSH_SKILL, 'w') as f: | |
| f.write(json_data) | |
| def load_json_via_sftp(): | |
| global sftp_matchmaker_client | |
| if not is_connected(): | |
| create_ssh_matchmaker_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| with sftp_matchmaker_client.open(SSH_SKILL, 'r') as f: | |
| data = json.load(f) | |
| ratings = [deserialize_rating(r) for r in data['ratings']] | |
| comparison_counts = np.array(data['comparison_counts']) | |
| total_comparisons = data['total_comparisons'] | |
| return ratings, comparison_counts, total_comparisons | |
| class RunningPivot(object): | |
| running_pivot = [] | |
| def matchmaker(num_players, k_group=4, not_run=[]): | |
| trueskill_env = TrueSkill() | |
| ratings, comparison_counts, total_comparisons = load_json_via_sftp() | |
| ratings = ratings[:num_players] | |
| comparison_counts = comparison_counts[:num_players, :num_players] | |
| # Randomly select a player | |
| # selected_player = np.random.randint(0, num_players) | |
| comparison_counts[RunningPivot.running_pivot, :] = float('inf') | |
| comparison_counts[not_run, :] = float('inf') | |
| selected_player = np.argmin(comparison_counts.sum(axis=1)) | |
| RunningPivot.running_pivot.append(selected_player) | |
| RunningPivot.running_pivot = RunningPivot.running_pivot[-5:] | |
| print(RunningPivot.running_pivot) | |
| selected_trueskill_score = trueskill_env.expose(ratings[selected_player]) | |
| trueskill_scores = np.array([trueskill_env.expose(p) for p in ratings]) | |
| trueskill_diff = np.abs(trueskill_scores - selected_trueskill_score) | |
| n = comparison_counts[selected_player] | |
| ucb_scores = ucb_score(trueskill_diff, total_comparisons, n) | |
| # Exclude self, select opponent with highest UCB score | |
| ucb_scores[selected_player] = -float('inf') | |
| ucb_scores[not_run] = -float('inf') | |
| opponents = np.argsort(ucb_scores)[-k_group + 1:].tolist() | |
| # Group players | |
| model_ids = [selected_player] + opponents | |
| random.shuffle(model_ids) | |
| return model_ids | |