Skip to content

Together AI

TogetherClient

Bases: BatchClient

Source code in dactyl_generation/together_generation.py
 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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
class TogetherClient(BatchClient):
    def __init__(self, api_key: str) -> None:
        """
        Constructor for Together AI client. 

        Args:
            api_key (str): API key
        """
        super().__init__()
        self.client = Together(api_key=api_key)

    def create_batch_job(self, jsonl_path: str) -> Dict[str, str]:
        """
        Creates batch job from JSONL file.

        Args:
            jsonl_path: Path to JSONL file.

        Returns:
            batch_job_input: Dictionary of batch job information.
        """
        file_resp = self.client.files.upload(
            file=jsonl_path, purpose="batch-api", check=False
        )
        prompts_df = pd.read_json(jsonl_path, lines=True)
        file_id = file_resp.id

        batch = self.client.batches.create(input_file_id=file_id, endpoint="/v1/chat/completions").job


        return {
            BATCH_ID: batch.id,
            INPUT_FILE: batch.input_file_id,
            CREATED: str(batch.created_at),
            MODEL: batch.x_model_id,
            PROMPTS: prompts_df.to_dict(orient="records")
        }

    def get_batch_job_output(self, batch_job_file: str) -> pd.DataFrame:
        """
        Retrieves batch job output.


        Args:
            batch_job_file: Path to batch job input, generated from the function `create_batch_job`.
        Returns:
            result_df: Pandas DataFrame of batch job output.
        """
        with open(batch_job_file, 'r') as f:
            data = json.load(f)

        batch_id = data[BATCH_ID]
        batch = self.client.batches.retrieve(batch_id)
        prompts_df = pd.DataFrame(data[PROMPTS])


        with self.client.files.with_streaming_response.content(id=batch.output_file_id) as response:
            results_df = pd.DataFrame([json.loads(line) for line in response.iter_lines()])

        prompts_df = prompts_df.merge(results_df, how="left", on=CUSTOM_ID)
        return prompts_df

__init__(api_key)

Constructor for Together AI client.

Parameters:

Name Type Description Default
api_key str

API key

required
Source code in dactyl_generation/together_generation.py
10
11
12
13
14
15
16
17
18
def __init__(self, api_key: str) -> None:
    """
    Constructor for Together AI client. 

    Args:
        api_key (str): API key
    """
    super().__init__()
    self.client = Together(api_key=api_key)

create_batch_job(jsonl_path)

Creates batch job from JSONL file.

Parameters:

Name Type Description Default
jsonl_path str

Path to JSONL file.

required

Returns:

Name Type Description
batch_job_input Dict[str, str]

Dictionary of batch job information.

Source code in dactyl_generation/together_generation.py
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def create_batch_job(self, jsonl_path: str) -> Dict[str, str]:
    """
    Creates batch job from JSONL file.

    Args:
        jsonl_path: Path to JSONL file.

    Returns:
        batch_job_input: Dictionary of batch job information.
    """
    file_resp = self.client.files.upload(
        file=jsonl_path, purpose="batch-api", check=False
    )
    prompts_df = pd.read_json(jsonl_path, lines=True)
    file_id = file_resp.id

    batch = self.client.batches.create(input_file_id=file_id, endpoint="/v1/chat/completions").job


    return {
        BATCH_ID: batch.id,
        INPUT_FILE: batch.input_file_id,
        CREATED: str(batch.created_at),
        MODEL: batch.x_model_id,
        PROMPTS: prompts_df.to_dict(orient="records")
    }

get_batch_job_output(batch_job_file)

Retrieves batch job output.

Parameters:

Name Type Description Default
batch_job_file str

Path to batch job input, generated from the function create_batch_job.

required

Returns: result_df: Pandas DataFrame of batch job output.

Source code in dactyl_generation/together_generation.py
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def get_batch_job_output(self, batch_job_file: str) -> pd.DataFrame:
    """
    Retrieves batch job output.


    Args:
        batch_job_file: Path to batch job input, generated from the function `create_batch_job`.
    Returns:
        result_df: Pandas DataFrame of batch job output.
    """
    with open(batch_job_file, 'r') as f:
        data = json.load(f)

    batch_id = data[BATCH_ID]
    batch = self.client.batches.retrieve(batch_id)
    prompts_df = pd.DataFrame(data[PROMPTS])


    with self.client.files.with_streaming_response.content(id=batch.output_file_id) as response:
        results_df = pd.DataFrame([json.loads(line) for line in response.iter_lines()])

    prompts_df = prompts_df.merge(results_df, how="left", on=CUSTOM_ID)
    return prompts_df