podcast-sponsor-remove

Attempt at identify sponsored segments in audio transcripts and removing them.
Log | Files | Refs

reformat.py (2712B)

      1 import sys
      2 import json
      3 from transformers import AutoTokenizer
      4 from tqdm import tqdm
      5 
      6 # IMPORTANT: The tokenizer MUST match the student model you are fine-tuning.
      7 STUDENT_MODEL_ID = "google/t5gemma-s-s-ul2-it"
      8 
      9 def reformat_with_template(input_filepath: str, output_filepath: str):
     10     """
     11     Reads a JSONL file with 'text' (prompt) and 'target' (completion) fields
     12     and reformats it into a single 'text' field using the model's chat template.
     13     """
     14     print(f"Loading tokenizer for {STUDENT_MODEL_ID}...")
     15     tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_ID)
     16     
     17     print(f"Reformatting {input_filepath}...")
     18     
     19     with open(input_filepath, 'r', encoding='utf-8') as infile, \
     20          open(output_filepath, 'w', encoding='utf-8') as outfile:
     21         
     22         # Count lines for tqdm progress bar
     23         num_lines = sum(1 for line in open(input_filepath, 'r', encoding='utf-8'))
     24         infile.seek(0) # Reset file pointer
     25         
     26         for line in tqdm(infile, total=num_lines, desc="Reformatting"):
     27             original_entry = json.loads(line)
     28             
     29             # The original prompt is in the 'text' field
     30             prompt_text = original_entry.get("text", "")
     31             # The verified completion is in the 'target' field
     32             completion_text = original_entry.get("target", "")
     33             
     34             # The prompt needs to end with the cue for the model to respond to
     35             if "Timestamp:" not in prompt_text:
     36                  prompt_text += "\n\nTimestamp:\n"
     37 
     38             # Create the message structure that the chat template expects
     39             messages = [
     40                 {"role": "user", "content": prompt_text},
     41                 {"role": "assistant", "content": completion_text}
     42             ]
     43             
     44             # Apply the template to create the new, single text field
     45             # add_generation_prompt=False is important for training data
     46             formatted_text = tokenizer.apply_chat_template(
     47                 messages, 
     48                 tokenize=False, 
     49                 add_generation_prompt=False
     50             )
     51             
     52             new_entry = {
     53                 "id": original_entry.get("id"),
     54                 "text": formatted_text,
     55                 "target": completion_text # Keep target for reference if needed
     56             }
     57             outfile.write(json.dumps(new_entry) + "\n")
     58 
     59     print(f"\nSuccessfully reformatted {num_lines} entries. Output saved to {output_filepath}.")
     60 
     61 if __name__ == "__main__":
     62     if len(sys.argv) < 3:
     63         print("Usage: python reformat_with_template.py <input_dataset.jsonl> <output_dataset_templated.jsonl>")
     64         sys.exit(1)
     65         
     66     reformat_with_template(sys.argv[1], sys.argv[2])