github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-05 18:14:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2025-04-18 06:08:04 +00:00
parent ad56d16c1d
commit d61f4a0a30
3 changed files with 43 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
applications/ColossalChat/coati/dataset/loader.py
View File

@ -438,6 +438,7 @@ class RawConversationDataset(Dataset):
self.tokenized_texts[index] = dict(tokens)
return self.tokenized_texts[index]
class AIMEDataset(Dataset):
"""
AIME dataset.
@ -458,12 +459,16 @@ class AIMEDataset(Dataset):
def __getitem__(self, index: int):
if self.tokenized_texts[index] is None:
message = self.raw_texts[index]
gt_answer = self.tokenizer.encode(message['answer'], padding="max_length", truncation=True, max_length=self.max_length, return_tensors="pt")
gt_answer = self.tokenizer.encode(
message["answer"],
padding="max_length",
truncation=True,
max_length=self.max_length,
return_tensors="pt",
)
def make_conv_hf(question):
msg = [
{"role": "user", "content": question}
]
msg = [{"role": "user", "content": question}]
return msg
message = make_conv_hf(message["question"])
@ -472,8 +477,14 @@ class AIMEDataset(Dataset):
self.tokenized_texts[index]["gt_answer"] = gt_answer.squeeze(1)
return self.tokenized_texts[index]
if __name__ == "__main__":
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("/home/share/data/model/Qwen2.5-3B")
dataset = AIMEDataset(tokenizer, "/home/yanglibing/workspace/PRIME/eval/data/AI-MO/aimo-validation-aime/aimo-validation-aime.jsonl", 512)
dataset = AIMEDataset(
tokenizer,
"/home/yanglibing/workspace/PRIME/eval/data/AI-MO/aimo-validation-aime/aimo-validation-aime.jsonl",
512,
)
print(dataset[0])

19
applications/ColossalChat/coati/distributed/producer.py
View File

@ -1,5 +1,4 @@
from collections import defaultdict
import os
from typing import Any, Dict, Optional
import numpy as np
@ -7,15 +6,12 @@ import ray
import ray.util.collective as cc
import torch
from coati.dataset.loader import RawConversationDataset
import wandb
from applications.ColossalChat.coati.distributed.reward.reward_fn import math_reward_fn
from coati.distributed.reward.verifiable_reward import VerifiableReward
from torch import nn
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoTokenizer
from applications.ColossalChat.build.lib.coati.models.utils import read_jsonl_file
from applications.ColossalChat.coati.dataset.loader import AIMEDataset
from applications.ColossalChat.coati.distributed.reward.reward_fn import math_reward_fn
from colossalai.utils import get_current_device
from .comm import ray_broadcast_tensor_dict
@ -192,7 +188,6 @@ class SimpleProducer(BaseProducer):
all_rewards = []
all_formats = []
all_accs = []
batch_reward_means = []
self.val_dataset = AIMEDataset(
tokenizer=self.tokenizer,
@ -229,11 +224,9 @@ class SimpleProducer(BaseProducer):
return {**tensors, **non_tensors}
self.val_dataloader = DataLoader(dataset=self.val_dataset,
batch_size=64,
shuffle=True,
drop_last=True,
collate_fn=collate_fn)
self.val_dataloader = DataLoader(
dataset=self.val_dataset, batch_size=64, shuffle=True, drop_last=True, collate_fn=collate_fn
)
all_rewards = torch.tensor([], device=self.device)
all_formats = torch.tensor([], device=self.device)
@ -249,7 +242,8 @@ class SimpleProducer(BaseProducer):
data = {k: v.view(-1, v.size(-1)) for k, v in test_output.items()}
# data = test_output
reward_group = self.reward_model(
data["input_ids"], gt_answer=data["gt_answer"], response_idx=data["response_idx"])
data["input_ids"], gt_answer=data["gt_answer"], response_idx=data["response_idx"]
)
rewards = torch.stack([x[0] for x in reward_group])
format_rewards = torch.stack([x[1] for x in reward_group])
@ -275,4 +269,3 @@ class SimpleProducer(BaseProducer):
f"acc={valid_metrics['avg_acc'].item():.4f}"
)
return valid_metrics