resolve comments

Author: gabriben

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
+from .data_utils import write_dataset
 from dynaconf import Dynaconf
 from .prompt_tones import get_all_tones, Tone
 import os

src/wraval/actions/action_inference.py

@@ -15,14 +15,14 @@ def run_inference(
     data_dir: str
 ) -> None:
     """Run inference on sentences using the specified model"""
-    d = load_latest_dataset(data_dir)
+    results = load_latest_dataset(data_dir)
 
-    if "rewrite" not in d.columns:
-        d["rewrite"] = None
-    if "inference_model" not in d.columns:
-        d["inference_model"] = None
+    if "rewrite" not in results.columns:
+        results["rewrite"] = None
+    if "inference_model" not in results.columns:
+        results["inference_model"] = None
 
-    tones = d["tone"].unique()
+    tones = results["tone"].unique()
     print(f"Found tones: {tones}")
 
     if settings.type != "all":
@@ -39,17 +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)
 
         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
+        new_results = pd.DataFrame({"synthetic_data" : queries, "tone" : tone})
+        new_results["rewrite"] = cleaned_output
+        new_results["inference_model"] = model_name
 
-        d = pd.concat([d, new], ignore_index=True)
+        results = pd.concat([results, new_results], ignore_index=True)
 
-    write_dataset(d, data_dir, "all", "csv")
+    write_dataset(results, data_dir, "all", "csv")

src/wraval/actions/completion.py

@@ -196,7 +196,6 @@ 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

@@ -9,16 +9,16 @@
 
 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)
+        write_dataset_s3(df,  bucket, prefix, file_prefix, format)
     else:
-        return write_dataset_local(df, data_dir, file_prefix, format)
+        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.csv")
         df.to_csv(temp_file, index=False)
@@ -28,7 +28,6 @@ def write_dataset_s3(
                            )
         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, format: str
@@ -42,8 +41,7 @@ def write_dataset_local(
         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/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)