Merge pull request #13 from amazon-science/agnostic-write-dataset

write_dataset that calls either s3 or local

Author: EetuaLaine

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
 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,19 +15,14 @@ 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
-
-    if "rewrite" not in d.columns:
-        d["rewrite"] = None
-    if "inference_model" not in d.columns:
-        d["inference_model"] = None
-
-    tones = d["tone"].unique()
+    results = load_latest_dataset(data_dir)
+
+    if "rewrite" not in results.columns:
+        results["rewrite"] = None
+    if "inference_model" not in results.columns:
+        results["inference_model"] = None
+
+    tones = results["tone"].unique()
     print(f"Found tones: {tones}")
 
     if settings.type != "all":
@@ -43,18 +39,17 @@ def run_inference(
 
         tone_prompt = get_prompt(Tone(tone))
 
-        queries = d[d["tone"] == tone]["synthetic_data"].unique()
+        queries = results[results["tone"] == tone]["synthetic_data"].unique()
 
         print(f"Processing {len(queries)} unique inputs for tone: {tone}")
 
         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_results = pd.DataFrame({"synthetic_data" : queries, "tone" : tone})
+        new_results["rewrite"] = cleaned_output
+        new_results["inference_model"] = model_name
+
+        results = pd.concat([results, new_results], ignore_index=True)
 
-    write_dataset_local(d, "./data", "all-tones")
-    if upload_s3:
-        write_dataset_to_s3(d, settings.s3_bucket, "inference/all", "csv")
+    write_dataset(results, 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/data_utils.py

@@ -7,30 +7,41 @@
 from urllib.parse import urlparse
 
 
-def write_dataset_to_s3(
-    df: pd.DataFrame, bucket: str, key_prefix: str, format: str
-) -> str:
+def write_dataset(
+    df: pd.DataFrame, data_dir: str, file_prefix: str, format: str
+):
+    if is_s3_path(data_dir):
+        bucket, prefix = parse_s3_path(data_dir)
+        write_dataset_s3(df,  bucket, prefix, file_prefix, format)
+    else:
+        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
+):
     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}")
-    return output_path
+    print(f"Saved locally to {output_path}")
 
 
 def add_timestamp_to_file_prefix(file_prefix, format):

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

src/wraval/model_artifacts/code/inference.py

@@ -27,4 +27,4 @@ def predict_fn(data, model_and_tokenizer, *args):
     model.config.pad_token_id = model.config.eos_token_id
     inputs = tokenizer(sentences, return_tensors="pt", padding=True).to(model.device)
     output_sequences = model.generate(**inputs, max_new_tokens=1024)
-    return tokenizer.batch_decode(output_sequences, skip_special_tokens=True)
+    return tokenizer.batch_decode(output_sequences)