import vllm
from vllm import LLM, SamplingParams
from vllm.sampling_params import GuidedDecodingParams
from pydantic import BaseModel
from judge_api import guess
LLM_BASE_URL = "https://openrouter.ai/api/v1"
LLM_API_KEY = "sk-or-v1-openrouter-api-key"
MODEL = 'openai/gpt-4.1'
import math, random
class GameClient:
def ndcg_at_10(self, predictions, correct_answer):
if correct_answer not in predictions:
return 0.0
try:
rank = predictions[:10].index(correct_answer) + 1
except ValueError:
return 0.0
return 1 / math.log2(rank + 1)
def hits_at_10(self, predictions, correct_answer):
return 1.0 if correct_answer in predictions[:10] else 0.0
def __init__(self):
self._random_options = ['gym', 'dinosaur', 'camel', 'desk', 'chicken', 'suitcase', 'thief', 'penguin', 'bat', 'painter', 'yogurt', 'chocolate', 'football', 'wallet', 'magician', 'shoes', 'bank', 'church', 'chewing gum', 'fashion', 'chainsaw', 'escalator', 'scarf', 'lawyer', 'eagle', 'credit card', 'garden hose', 'glider', 'crosswalk', 'subway', 'fireworks', 'marshmallow', 'cookies', 'curtains', 'dining room', 'cars', 'wedding', 'guitar', 'coffee', 'mouse', 'meat', 'scale', 'train tracks', 'zebra', 'fairy', 'quit', 'museum', 'kangaroo', 'surfboard', 'cheese', 'nightmare', 'jellyfish', 'koala', 'strawberry', 'tiger', 'mailbox', 'kettle', 'potato', 'janitor', 'lighthouse', 'crocodile', 'charger', 'doctor', 'peacock', 'peanut', 'popcorn', 't-shirt', 'fertilizer', 'keyboard', 'umbrella', 'pool', 'watercolor', 'mango', 'xylophone', 'bathroom', 'ice cube', 'giraffe', 'garage', 'cabin', 'plankton', 'pig', 'vulture', 'frame', 'polar bear', 'microscope', 'snake', 'skeleton', 'rocket', 'backpack', 'jacket', 'bedroom', 'castle', 'horse', 'dragonfly', 'hotel', 'cyclist', 'mask', 'restaurant', 'toothpaste', 'angel', 'whistle', 'wrestling', 'eclipse', 'hermit crabs', 'horn', 'boxers', 'volcano', 'fire station', 'toothbrush', 'egg', 'straw', 'rice', 'diamond', 'vitamins', 'tricycle', 'bottle-opener', 'panther', 'ice skates', 'theater', 'gas mask', 'game console', 'path', 'scorpion', 'snowboard', 'crab', 'pie', 'octopus', 'mustache', 'pepper grinder', 'swings', 'palm tree', 'well', 'sewing machine', 'key', 'station', 'mosque', 'chameleon', 'cherry', 'parrot', 'leggings', 'radio', 'brick', 'sunflower', 'hammer', 'carrot', 'radar', 'kite', 'bathtub', 'rhinoceros', 'spoon', 'orchestra', 'gravity', 'flute', 'lipstick', 'school', 'meteorite', 'politician', 'ladder', 'lawnmower', 'computer', 'wheel', 'airport', 'firefighter', 'porch', 'police station', 'queen', 'mayonnaise', 'alumunium foil', 'lion', 'helmet', 'teacher', 'tea', 'fan', 'piano', 'snail', 'farmer', 'harbor', 'nurse', 'sunglasses', 'bee', 'postal worker', 'market', 'plank', 'steering wheel', 'squirrel', 'netting', 'dragon', 'cafeteria', 'millennium', 'spinach', 'fork', 'cabbage', 'ping-pong', 'lock', 'submarine', 'dictionary', 'vaccine', 'soda', 'skirt', 'toaster', 'shorts', 'circus', 'flowerpot', 'lobster', 'rainbow', 'cockroach', 'frog', 'basket ball', 'chilli pepper', 'pajamas', 'crossword', 'light bulb', 'drill', 'beaver', 'daisy', 'river', 'yo-yo', 'harmonica', 'soap', 'igloo', 'sausage', 'deer', 'sailboat', 'fish', 'mosquito', 'can', 'rat', 'frying pan', 'barcode', 'sunscreen', 'ferret', 'whale', 'duck', 'shirt', 'vacuum', 'detective', 'perfume', 'seal', 'raincoat', 'alien', 'bull', 'nest', 'butterfly', 'eraser', 'hedgehog', 'panda', 'refrigerator', 'monocle', 'window', 'kitchen', 'mole', 'speaker', 'waiter', 'salad', 'dolphin', 'storm', 'drums', 'spiderweb', 'bicycle', 'monkey', 'flamingo', 'prison', 'bowling', 'pencil sharpner', 'photo', 'printer', 'robe', 'seahorse', 'doorbell', 'gloves', 'alcohol', 'diving suit', 'shotgun', 'hairbrush', 'cactus', 'ambulance', 'hula hoop', 'snowman', 'mountain', 'unicorn', 'suit', 'cake', 'cow', 'sled', 'boar', 'barbecue', 'trash can', 'slingshot', 'banana', 'dam', 'hat', 'milk', 'shell', 'broom', 'fisherman', 'bucket', 'bell', 'tracktor', 'fly', 'spider', 'carpet', 'coconut tree', 'movie theater', 'socks', 'soldier', 'watering can', 'accountant', 'microphone', 'toothpick', 'wolf', 'trumpet', 'apple', 'library', 'cork', 'zipper', 'pan', 'doghouse', 'dynamite', 'swan', 'grasshopper', 'beach', 'starfish', 'police officer', 'board game', 'magnet', 'cucumber', 'fire extinguisher', 'sundial', 'mechanic', 'lighter', 'shovel', 'shark', 'notebook', 'ostrich', 'bodyguard', 'binoculars', 'parachute', 'drone', 'kiwi', 'ghost', 'baker', 'robot', 'postcard', 'horseshoe', 'karaoke', 'billiards', 'palace', 'hospital', 'compass', 'truck', 'holiday', 'lake', 'cave', 'space station', 'mushroom', 'magnifying glass', 'fox', 'bread', 'rose', 'windmill', 'pirate', 'earring', 'hunter', 'princess', 'calculator', 'clown', 'watch', 'pilot', 'mustard', 'swordfish', 'darts', 'microwave oven', 'plumber', 'sword']
def _make_api_call(self, clues, options):
return guess(clues, options)
def simulate_game(self, clues, expected_answer, distractors = []):
expected_answer = expected_answer.lower()
if expected_answer in distractors:
options = []
else:
options = [expected_answer]
if len(distractors) > 0:
options.extend([d.lower() for d in distractors])
options = options[:100]
if len(options) < 100:
random.seed(expected_answer)
options.extend(random.choices(self._random_options, k=100-len(options)))
random.shuffle(options)
try:
predictions = self._make_api_call(clues, options)
return {
"predictions": predictions,
"hit@10": self.hits_at_10(predictions, expected_answer),
"NDCG@10": self.ndcg_at_10(predictions, expected_answer)
}
except Exception as e:
return {
"predictions": [],
"hit@10": 0.0,
"NDCG@10": 0.0
}
from datasets import load_dataset
DESCRIPTIONS = "IOAI-official/ioai2025-onsite-concepts-hint-descriptions"
hint_descriptions = load_dataset(DESCRIPTIONS)['train']
hint_descriptions = {
x['ID']: {'description': x['Description'], 'icons': x['image']}
for x in hint_descriptions
}
def tns_desc(description: str):
return description.replace("\n", "-")
valid_hints = [tns_desc(x['description']) for x in hint_descriptions.values()]
hint_to_id = {
tns_desc(x['description']): xid
for xid, x in hint_descriptions.items()
}
print(valid_hints)
print(hint_to_id)
TRAINING_SET = "IOAI-official/ioai2025-onsite-concepts-train"
dev = load_dataset(TRAINING_SET)['train']
print(dev)
print(dev[0])
Fine Tuning
from openai import AsyncOpenAI
oai_client = AsyncOpenAI(
base_url=LLM_BASE_URL,
api_key=LLM_API_KEY
)
from unsloth import FastLanguageModel
import torch
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "unsloth/Qwen3-0.6B",
max_seq_length = 4096,
full_finetuning = False,
)
from typing import Literal, List, Optional
def get_prompt(valid_hints, options, answer):
hints_str = str(valid_hints).replace("'", '"')
user_content = (
f"Your valid list of clues are: {hints_str}. "
"Your job is a clue giver. You will help the guesser pick the correct answer from a range of options. "
"You must output AT LEAST 1 hint and AT MOST 8 hints for each list of hints. "
"Please output a json object with the key 'hints_1', 'hints_2', 'hints_3', and 'hints_4', and the value being a list of hints. "
"IMPORTANT: You must ONLY use hint names from the valid hints list provided above. Do not create new hint names. "
'Example for the answer "knight": '
'{"hints_1":["Human-Historical-Real","Attack-Conflict - Combat-Weapon","Masculine-Male"], '
'"hints_2":["Character-Fictional-Imaginary","Work-Occupation","Defence-Protection-Wall","Arm-Hand-Finger"], '
'"hints_3":["Clothing-Accessory-Cotume","Power-Politics","Building-Monument-City"], '
'"hints_4":["Life-Heart-Love","Metal","Torso - Back-Stomach-Body"]} '
f"The answer word is :{answer}. Please construct sequences of hints that are most relevant to it. Remember to adhere to your instructions."
)
prompt = [
{"role": "system", "content": "You are an AI giving hints to another AI in a concept guessing game. Generate up to 4 list of concept markers, or hints, to help another AI guess the correct answer."},
{"role": "user", "content": user_content}
]
text = tokenizer.apply_chat_template(
prompt,
tokenize=False,
add_generation_prompt=True,
)
return text
ValidHint = Literal[*valid_hints]
class Hints(BaseModel):
hints_1: List[ValidHint]
hints_2: List[ValidHint]
hints_3: List[ValidHint]
hints_4: List[ValidHint]
def to_result(self):
hints = [self.hints_1, self.hints_2, self.hints_3, self.hints_4]
result = []
for hintlist in hints:
cur_hintlist = [hint_to_id[hint] for hint in hintlist[:8]]
result.append(cur_hintlist)
return result
async def gen_dat(label, options):
messages = [{
"role": "user",
"content": get_prompt(valid_hints, options, label)
}]
response = await oai_client.beta.chat.completions.parse(
model=MODEL,
messages=messages,
response_format=Hints
)
messages.append(
{
"role": "assistant",
"content": response.choices[0].message.parsed.model_dump_json()
}
)
return messages
sample_res = await gen_dat(dev[0]['label'], dev[0]['options'])
print(sample_res[-1]['content'])
import random
all_options = list(set([opt for x in dev for opt in x['options']]))
print(len(all_options))
def gen_inout():
options = random.sample(all_options, 100)
label = random.choice(options)
return (label, options)
initial_corpos = [gen_inout() for _ in range(1000)]
print(initial_corpos[0])
print(await gen_dat(*initial_corpos[0]))
from asyncio import Semaphore
from tqdm.asyncio import tqdm_asyncio
sema = Semaphore(90)
async def gen_with_sema(inp):
async with sema:
return await gen_dat(*inp)
tasks = [gen_with_sema(inp) for inp in initial_corpos]
results = await tqdm_asyncio.gather(*tasks, total=len(tasks))
len(results)
SAVE_DIR = "out"
import pickle
import os
os.makedirs(SAVE_DIR, exist_ok=True)
with open(f"{SAVE_DIR}/ft_data.pickle", "wb") as f:
pickle.dump(results, f)
import pickle
with open(f"{SAVE_DIR}/ft_data.pickle", "rb") as f:
results = pickle.load(f)
print(f"Loaded {len(results)} results.")
formatted_results_qwen = [[item['content'] for item in result] for result in results]
results = formatted_results_qwen
results[0]
!nvidia-smi
model = FastLanguageModel.get_peft_model(
model,
r = 32,
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",],
lora_alpha = 32,
lora_dropout = 0,
bias = "none",
use_gradient_checkpointing = "unsloth",
random_state = 3407,
use_rslora = False,
loftq_config = None,
)
processed_conversations = results
processed_conversations[0]
import pandas as pd
from datasets import Dataset
training_set = pd.Series(processed_conversations)
training_set.name = "text"
training_dataset = Dataset.from_pandas(pd.DataFrame(training_set))
from trl import SFTTrainer, SFTConfig
trainer = SFTTrainer(
model = model,
tokenizer = tokenizer,
train_dataset = training_dataset,
eval_dataset = None,
args = SFTConfig(
dataset_text_field = "text",
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
warmup_steps = 5,
num_train_epochs = 1,
max_steps = 100,
learning_rate = 2e-4,
optim = "adamw_8bit",
weight_decay = 0.01,
lr_scheduler_type = "linear",
seed = 3407,
logging_steps=1,
logging_dir="./logs",
logging_first_step=True,
report_to="none",
disable_tqdm=True,
log_level="debug"
),
)
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")
import torch
def print_gpu_utilization():
allocated = torch.cuda.memory_allocated()/1e9
reserved = torch.cuda.memory_reserved()/1e9
print(f"显存使用: {allocated:.2f}GB / {reserved:.2f}GB")
print_gpu_utilization()
trainer_stats = trainer.train()
print_gpu_utilization()
used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory / max_memory * 100, 3)
lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(
f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training."
)
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
model.save_pretrained_merged(f"{SAVE_DIR}/model", tokenizer, save_method = "merged_16bit",)
torch.cuda.empty_cache()
class ClueGiver:
def __init__(self):
self.llm = LLM(
f"{SAVE_DIR}/model",
gpu_memory_utilization=0.8,
max_model_len=4096,
tensor_parallel_size=1
)
json_schema = Hints.model_json_schema()
self.sampling_params = SamplingParams(
guided_decoding=GuidedDecodingParams(
json=json_schema,
),
max_tokens=5096,
frequency_penalty=0.5,
presence_penalty=0.8
)
def construct_clues(self, answers: List[str], options: List[List[str]]):
prompts = []
for answer, options in zip(answers, options):
prompt = get_prompt(valid_hints, options, answer)
prompts.append(prompt)
hints_batch = self.llm.generate(prompts=prompts, sampling_params=self.sampling_params)
results = []
for i, hints in enumerate(hints_batch):
txt = hints.outputs[0].text
print(txt)
try:
hints_obj = Hints.model_validate_json(txt)
results.append(hints_obj.to_result())
except Exception as e:
print(f"Error parsing hints for answer {answers[i]}: {e}")
results.append([[1,2,3,4]])
return results
torch.cuda.empty_cache()
import gc
gc.collect()
clue_giver = ClueGiver()
res_clues = clue_giver.construct_clues([x['label'] for x in dev], [x['options'] for x in dev])
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed
def simulate_one(i_data):
i, data = i_data
clues = res_clues[i]
prediction = game_client.simulate_game(clues, data['label'])
return prediction
game_client = GameClient()
predictions = []
with ThreadPoolExecutor() as executor:
futures = [executor.submit(simulate_one, (i, data)) for i, data in enumerate(dev)]
for f in tqdm(as_completed(futures), total=len(futures)):
predictions.append(f.result())
print("Final Score: ")
print(sum([p['hit@10'] for p in predictions]) * 0.9 + sum([p['NDCG@10'] for p in predictions]) * 0.1)
Submission
VALID_PATH = "IOAI-official/ioai2025-onsite-concepts-validation"
TEST_PATH = "IOAI-official/ioai2025-onsite-concepts-test"
testset_a = load_dataset(VALID_PATH)["test"]
testset_b = load_dataset(TEST_PATH)["test"]
clues_a = clue_giver.construct_clues([x['label'] for x in testset_a], [x['options'] for x in testset_a])
clues_b = clue_giver.construct_clues([x['label'] for x in testset_b], [x['options'] for x in testset_b])
import json
OUTPUT_DIR = "out"
def write_clues(clues: List[List[List[int]]], path: str):
with open(path, 'w') as f:
for c in clues:
f.write(json.dumps(c) + '\n')
os.makedirs(OUTPUT_DIR, exist_ok=True)
write_clues(clues_a, OUTPUT_DIR + "/clues_a.jsonl")
write_clues(clues_b, OUTPUT_DIR + "/clues_b.jsonl")
del clue_giver
torch.cuda.empty_cache()
import gc
gc.collect()