mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-19 12:57:58 +00:00
233 lines
No EOL
10 KiB
Python
233 lines
No EOL
10 KiB
Python
import os
|
|
import random
|
|
import sys
|
|
import traceback
|
|
import csv
|
|
from typing import List, Optional, Tuple, Any, Dict
|
|
|
|
from datasets import load_dataset
|
|
from pydantic import Field
|
|
|
|
from atroposlib.envs.base import BaseEnv, BaseEnvConfig, OpenaiConfig, ScoredDataGroup
|
|
from atroposlib.type_definitions import GameHistory, Item
|
|
from atroposlib.utils.tokenize_for_trainer import tokenize_for_trainer
|
|
|
|
|
|
class UFCEnvConfig(BaseEnvConfig):
|
|
"""Configuration for the UFC Environment"""
|
|
fighter_stats_path: str = Field(os.path.join(os.path.dirname(__file__), "fighter_stats.csv"), description="Path to fighter stats CSV")
|
|
fight_data_path: str = Field(os.path.join(os.path.dirname(__file__), "large_dataset.csv"), description="Path to large fight dataset CSV")
|
|
max_steps: int = Field(1, description="Only one step per fight prediction")
|
|
temperature: float = Field(0.7, description="Temperature for generation diversity")
|
|
top_p: float = Field(0.95, description="Top p for nucleus sampling")
|
|
|
|
|
|
class UFCEnv(BaseEnv):
|
|
"""UFC Fight Prediction Environment"""
|
|
name = "ufc_predictor"
|
|
env_config_cls = UFCEnvConfig
|
|
|
|
def __init__(self, config: UFCEnvConfig, server_configs: List[OpenaiConfig], slurm=True, testing=False):
|
|
super().__init__(config, server_configs, slurm, testing)
|
|
self.fighter_stats = {}
|
|
self.fight_data = []
|
|
self.current_index = 0
|
|
self.inference_server = self.server.servers[0] # Get first server as inference server
|
|
|
|
async def setup(self):
|
|
"""Load the fighter stats and fight data"""
|
|
try:
|
|
print("Loading fighter stats from:", self.config.fighter_stats_path)
|
|
with open(self.config.fighter_stats_path, encoding="utf-8") as f:
|
|
reader = csv.DictReader(f)
|
|
self.fighter_stats = {row["name"]: row for row in reader}
|
|
print(f"Loaded stats for {len(self.fighter_stats)} fighters")
|
|
|
|
print("Loading fight data from:", self.config.fight_data_path)
|
|
with open(self.config.fight_data_path, encoding="utf-8") as f:
|
|
reader = csv.DictReader(f)
|
|
self.fight_data = list(reader)
|
|
print(f"Loaded {len(self.fight_data)} fights")
|
|
|
|
except Exception as e:
|
|
print(f"Error loading data: {e}")
|
|
traceback.print_exc()
|
|
sys.exit(1)
|
|
|
|
async def get_next_item(self) -> Optional[Item]:
|
|
"""Get the next fight from the dataset"""
|
|
try:
|
|
if self.current_index >= len(self.fight_data):
|
|
return None
|
|
fight = self.fight_data[self.current_index]
|
|
self.current_index += 1
|
|
|
|
r_fighter = fight["r_fighter"]
|
|
b_fighter = fight["b_fighter"]
|
|
r_stats = self.fighter_stats.get(r_fighter, {})
|
|
b_stats = self.fighter_stats.get(b_fighter, {})
|
|
|
|
# Format the prompt
|
|
def stats_str(name, stats):
|
|
if not stats:
|
|
return f"{name}: (No stats available)"
|
|
return (
|
|
f"Name: {name}\n"
|
|
f"Wins: {stats.get('wins','?')} Losses: {stats.get('losses','?')} Age: {stats.get('age','?')}\n"
|
|
f"Height: {stats.get('height','?')} cm Weight: {stats.get('weight','?')} kg Reach: {stats.get('reach','?')} cm Stance: {stats.get('stance','?')}\n"
|
|
f"SLpM: {stats.get('SLpM','?')} Sig Str Acc: {stats.get('sig_str_acc','?')} SApM: {stats.get('SApM','?')} Str Def: {stats.get('str_def','?')}\n"
|
|
f"TD Avg: {stats.get('td_avg','?')} TD Acc: {stats.get('td_acc','?')} TD Def: {stats.get('td_def','?')} Sub Avg: {stats.get('sub_avg','?')}\n"
|
|
)
|
|
|
|
prompt_text = (
|
|
"🎤 LADIES AND GENTLEMEN! Welcome to the most electrifying show in sports entertainment - the UFC Fight Prediction Show! "
|
|
"Let's break down this matchup that's got everyone talking!\n\n"
|
|
f"*Drumroll please* In the red corner, we have :\n{stats_str(r_fighter, r_stats)}\n\n"
|
|
f"And in the blue corner:\n{stats_str(b_fighter, b_stats)}\n\n"
|
|
"Now, as your favorite fight analyst who's definitely not just making this up as I go along, I want you to:\n"
|
|
"1. Break down these fighters like you're explaining why your favorite TV show character would win in a fight\n"
|
|
"2. Compare their styles\n"
|
|
"3. Point out their advantages\n"
|
|
"Give us your best fight commentary! Make it exciting, make it dramatic, make it sound like you're calling the fight live! "
|
|
"Throw in some classic commentator phrases, maybe a 'OH MY GOODNESS!' or two, and definitely some dramatic pauses for effect.\n\n"
|
|
"End your masterpiece with the winner's name in this exact format:\n"
|
|
"\\boxed{fighter name}"
|
|
)
|
|
|
|
prompt = tuple([
|
|
frozenset({"role": "user", "content": prompt_text}.items())
|
|
])
|
|
|
|
winner = fight.get("winner", "") # Red or Blue
|
|
winner_name = r_fighter if winner == "Red" else b_fighter if winner == "Blue" else ""
|
|
ground_truth = f"Answer: {winner_name}" if winner_name else ""
|
|
|
|
return (prompt, ground_truth, None)
|
|
|
|
except Exception as e:
|
|
print(f"Error in get_next_item: {e}")
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
async def collect_trajectories(self, item: Item) -> Tuple[List[Tuple[GameHistory, str, Optional[str]]], List[Item]]:
|
|
to_score = []
|
|
to_backlog = []
|
|
|
|
system_msg = {
|
|
"role": "system",
|
|
"content": (
|
|
"You are an expert MMA analyst. You will be given two UFC fighters and their stats. "
|
|
"Your task is to predict the winner of the fight based on their statistics.\n\n"
|
|
"IMPORTANT: You MUST format your response in exactly two parts:\n"
|
|
"1. First, analyze the fighters' stats and explain create a fight commentary\n"
|
|
"2. Then on a new line, give ONLY your final prediction in this exact format:\n"
|
|
"\\boxed{fighter name}\n\n"
|
|
"For example:\n"
|
|
"After analyzing stats... [your analysis here]\n"
|
|
"\\boxed{John Smith}\n\n"
|
|
"If you do not end your response with the \\boxed{} format, you will receive a score of -1.0."
|
|
)
|
|
}
|
|
|
|
user_msg = {
|
|
"role": "user",
|
|
"content": dict(item[0][0])["content"]
|
|
}
|
|
|
|
messages = [system_msg, user_msg]
|
|
|
|
try:
|
|
chat_completions = await self.inference_server.chat_completion(
|
|
messages=messages,
|
|
n=self.config.group_size,
|
|
max_tokens=2048, # Increased from 512 to allow for longer, more detailed fight commentaries
|
|
temperature=self.config.temperature,
|
|
top_p=self.config.top_p,
|
|
timeout=60,
|
|
)
|
|
for choice in chat_completions.choices:
|
|
assistant_msg = {"role": "assistant", "content": choice.message.content}
|
|
history = [
|
|
{"role": "system", "content": system_msg["content"]},
|
|
{"role": "user", "content": user_msg["content"]},
|
|
{"role": "assistant", "content": choice.message.content}
|
|
]
|
|
to_score.append((history, item[1], None))
|
|
except Exception as e:
|
|
print(f"Error in collect_trajectories: {e}")
|
|
traceback.print_exc()
|
|
to_backlog.append(item)
|
|
|
|
if not to_score:
|
|
return None, to_backlog
|
|
|
|
scored_data = await self.score(to_score)
|
|
return scored_data, to_backlog
|
|
|
|
async def score(self, rollout_group_data) -> Optional[ScoredDataGroup]:
|
|
if not rollout_group_data:
|
|
return None
|
|
|
|
scores = ScoredDataGroup()
|
|
scores["tokens"] = []
|
|
scores["masks"] = []
|
|
scores["scores"] = []
|
|
scores["advantages"] = None
|
|
scores["ref_logprobs"] = None
|
|
scores["messages"] = None
|
|
scores["group_overrides"] = {"group_size": self.config.group_size}
|
|
scores["overrides"] = None
|
|
scores["ground_truths"] = []
|
|
|
|
random.shuffle(rollout_group_data)
|
|
for item in rollout_group_data:
|
|
out = tokenize_for_trainer(self.tokenizer, item[0])
|
|
tokens = out["tokens"]
|
|
masks = out["masks"]
|
|
|
|
try:
|
|
# Extract prediction and ground truth
|
|
reply = item[0][-1]["content"]
|
|
ground_truth = item[1].strip().lower()
|
|
|
|
# Extract name from ground truth (format: "answer: name")
|
|
ground_truth_name = ground_truth.replace("answer:", "").strip()
|
|
|
|
# Extract name from \boxed{name} format
|
|
import re
|
|
boxed_match = re.search(r"\\boxed{([^}]+)}", reply)
|
|
if boxed_match:
|
|
prediction = boxed_match.group(1).strip().lower()
|
|
# Compare just the names
|
|
reward = 1.0 if prediction == ground_truth_name else -1.0
|
|
else:
|
|
# No boxed answer found
|
|
reward = -1.0
|
|
|
|
except Exception as e:
|
|
print(f"Error scoring response: {e}")
|
|
reward = -1.0
|
|
ground_truth = item[1] if isinstance(item[1], str) else ""
|
|
|
|
if len([i for i in masks if i != -100]) < 10:
|
|
continue
|
|
|
|
scores["tokens"].append(tokens)
|
|
scores["masks"].append(masks)
|
|
scores["scores"].append(reward)
|
|
scores["ground_truths"].append(ground_truth)
|
|
|
|
if len(scores["tokens"]) >= self.config.group_size:
|
|
break
|
|
|
|
if not scores["tokens"]:
|
|
return None
|
|
|
|
return scores
|
|
|
|
async def evaluate(self, *args, **kwargs):
|
|
"""No-op evaluation"""
|
|
return
|
|
|
|
if __name__ == "__main__":
|
|
UFCEnv.cli() |