File size: 1,052 Bytes
0c51b93
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import json

from sotopia.database.logs import BaseEpisodeLog


class FakeEpisodeLog(BaseEpisodeLog):
    pk: str


def jsonl_to_episodes(
    jsonl_file_path: str,
) -> list[FakeEpisodeLog]:
    """Load episodes from a jsonl file.

    Args:
        jsonl_file_path (str): The file path.

    Returns:
        list[FakeEpisodeLog]: List of episodes that fakes an EpisodeLog object.
    """
    episodes = []
    with open(jsonl_file_path, "r") as f:
        for line in f:
            data = json.loads(line)
            episode = FakeEpisodeLog(
                pk=data["episode_id"],
                environment=data["environment_id"],
                agents=data["agent_ids"],
                tag=data["experiment_tag"],
                models=data["experiment_model_name_pairs"],
                messages=data["raw_messages"],
                reasoning=data["reasoning"],
                rewards=data["raw_rewards"],
                rewards_prompt=data["raw_rewards_prompt"],
            )
            episodes.append(episode)
    return episodes