write_dataset that calls either s3 or local

Author: gabriben

config/settings.toml

@@ -1,6 +1,7 @@
 [default]
 region = 'us-east-1'
-data_dir = 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
+data_dir = "./data" 
+# 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
 deploy_bucket_name = 'llm-finetune-us-east-1-{aws_account}'
 deploy_bucket_prefix = 'models'
 sagemaker_execution_role_arn = 'arn:aws:iam::{aws_account}:role/sagemaker-execution-role-us-east-1'

pyproject.toml

@@ -46,5 +46,4 @@ where = ["src"]
 
 [tool.setuptools.data-files]
 # This copies the config file into the installation (non-package dir)
-"config" = ["config/settings.toml"]
-
+"config" = ["config/settings.toml"]

src/wraval/actions/action_generate.py

@@ -3,7 +3,7 @@
 # // SPDX-License-Identifier: Apache-2.0
 #
 import pandas as pd
-from .data_utils import write_dataset_local, write_dataset_to_s3
+from .data_utils import write_dataset_local, write_dataset
 from dynaconf import Dynaconf
 from .prompt_tones import get_all_tones, Tone
 import os
@@ -79,6 +79,4 @@ def generate_tone_data(
 
     combined = pd.concat(datasets, ignore_index=True)
 
-    write_dataset_local(combined, settings.data_dir, "all-tones")
-    if upload_s3:
-        write_dataset_to_s3(combined, settings.s3_bucket, "generate/all", "csv")
+    write_dataset(combined, settings.data_dir, "all", "csv")

src/wraval/actions/action_inference.py

@@ -2,8 +2,9 @@
 # // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # // SPDX-License-Identifier: Apache-2.0
 #
+import pandas as pd
 from dynaconf import Dynaconf
-from .data_utils import write_dataset_local, write_dataset_to_s3, load_latest_dataset
+from .data_utils import write_dataset, load_latest_dataset
 from .prompt_tones import get_prompt, Tone
 from .model_router import route_completion
 
@@ -14,12 +15,7 @@ def run_inference(
     data_dir: str
 ) -> None:
     """Run inference on sentences using the specified model"""
-    try:
-        d = load_latest_dataset(data_dir)
-        print(f"Loaded dataset with {len(d)} rows")
-    except FileNotFoundError:
-        print("No dataset found. Please generate data first.")
-        return
+    d = load_latest_dataset(data_dir)
 
     if "rewrite" not in d.columns:
         d["rewrite"] = None
@@ -49,12 +45,11 @@ def run_inference(
 
         outputs = route_completion(settings, queries, tone_prompt)
 
-        for query, output in zip(queries, outputs):
-            mask = (d["synthetic_data"] == query) & (d["tone"] == tone)
-            cleaned_output = output.strip().strip('"')
-            d.loc[mask, "rewrite"] = cleaned_output
-            d.loc[mask, "inference_model"] = model_name
+        cleaned_output = [o.strip().strip('"') for o in outputs]
+        new = pd.DataFrame({"synthetic_data" : queries, "tone" : tone})
+        new["rewrite"] = cleaned_output
+        new["inference_model"] = model_name
 
-    write_dataset_local(d, "./data", "all-tones")
-    if upload_s3:
-        write_dataset_to_s3(d, settings.s3_bucket, "inference/all", "csv")
+        d = pd.concat([d, new], ignore_index=True)
+
+    write_dataset(d, data_dir, "all", "csv")

src/wraval/actions/action_llm_judge.py

@@ -5,7 +5,7 @@
 import pandas as pd
 from typing import List, Dict, Optional
 from dynaconf import Dynaconf
-from .data_utils import write_dataset_local, write_dataset_to_s3, load_latest_dataset
+from .data_utils import write_dataset, load_latest_dataset
 from .prompts_judge import generate_input_prompt, generate_system_prompt, get_rubric, rewrite_prompt
 
 from .completion import batch_get_bedrock_completions
@@ -144,9 +144,7 @@ def judge(
         d.loc[mask, dmt.columns] = dmt.values
 
     # Save results
-    write_dataset_local(d, "./data", "all-tones")
-    if upload_s3:
-        write_dataset_to_s3(d, settings.s3_bucket, "inference/all", "csv")
+    write_dataset(d, settings.data_dir, "all-tones", "csv")
 
 def rewrite_judge(
     model_id: str,

src/wraval/actions/completion.py

@@ -196,6 +196,7 @@ def invoke_sagemaker_endpoint(
             Body=input_string.encode("utf-8"),
             ContentType="application/json",
         )
+        import pdb; pdb.set_trace()
         json_output = response["Body"].readlines()
         plain_output = "\n".join(json.loads(json_output[0]))
         last_assistant = extract_last_assistant_response(plain_output)

src/wraval/actions/data_utils.py

@@ -7,26 +7,39 @@
 from urllib.parse import urlparse
 
 
-def write_dataset_to_s3(
-    df: pd.DataFrame, bucket: str, key_prefix: str, format: str
+def write_dataset(
+    df: pd.DataFrame, data_dir: str, file_prefix: str, format: str
+) -> str:
+    if is_s3_path(data_dir):
+        bucket, prefix = parse_s3_path(data_dir)
+        return write_dataset_s3(df,  bucket, prefix, file_prefix, format)
+    else:
+        return write_dataset_local(df, data_dir, file_prefix, format)
+
+def write_dataset_s3(
+    df: pd.DataFrame, bucket: str, prefix: str, file_prefix: str, format: str
 ) -> str:
     with tempfile.TemporaryDirectory() as temp_dir:
-        temp_file = os.path.join(temp_dir, "temp.jsonl")
-        df.to_json(temp_file, orient="records", lines=bool(format == "jsonl"))
+        temp_file = os.path.join(temp_dir, "temp.csv")
+        df.to_csv(temp_file, index=False)
         s3_client = boto3.client("s3")
-        key = add_timestamp_to_file_prefix(key_prefix, format)
-        print(f"Writing dataset to bucket {bucket} and key {key}.")
+        key = os.path.join(prefix, 
+                           add_timestamp_to_file_prefix(file_prefix, format)
+                           )
+        print(f"Writing dataset to s3://{bucket}/{key}")
         s3_client.upload_file(temp_file, bucket, key)
     return f"s3://{bucket}/{key}"
 
-
-def write_dataset_local(df: pd.DataFrame, data_dir: str, file_prefix: str) -> str:
+def write_dataset_local(
+    df: pd.DataFrame, data_dir: str, file_prefix: str, format: str
+    ) -> str:
     # Expand home directory and create if needed
     data_dir = os.path.expanduser(data_dir)
     os.makedirs(data_dir, exist_ok=True)
 
     output_path = os.path.join(
-        data_dir, add_timestamp_to_file_prefix(file_prefix, "csv")
+        data_dir, 
+        add_timestamp_to_file_prefix(file_prefix, format)
     )
     df.to_csv(output_path, index=False)
     print(f"Saved to {output_path}")

src/wraval/actions/model_router.py

@@ -5,7 +5,7 @@
 )
 from .format import format_prompt
 from transformers import AutoTokenizer
-from tqdm import tqdm
+from tqdm.auto import tqdm
 from typing import List
 from dynaconf.base import LazySettings
 from abc import ABC, abstractmethod