6 files changed, 186 insertions, 25 deletions

diff --git a/.gitignore b/.gitignore
index 6ade941..252715e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,7 +1,7 @@
 venv
 *~
-.*
 *.mp3
 *.wav
 *.json
 *.jsonl
+*/
diff --git a/finetune.py b/finetune.py
index bca7ac1..05b558f 100644
--- a/finetune.py
+++ b/finetune.py
@@ -1,15 +1,17 @@
 import torch
 from datasets import load_dataset
-from peft import LoraConfig
-from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from trl import SFTTrainer, SFTConfig
 import os
 
-#MODEL_ID = "google/gemma-3-270m-it"
-MODEL_ID = "google/gemma-3-1b-it"
+#MODEL_ID = "google/flan-t5-large"
+#MODEL_ID = "google/t5gemma-b-b-prefixlm-it"
+#MODEL_ID = "google/long-t5-tglobal-base"
+MODEL_ID = "google/t5gemma-s-s-ul2-it"
 DATASET_TRAIN_PATH = "train.jsonl"
 DATASET_VAL_PATH = "validation.jsonl"
-OUTPUT_DIR = "./gemma3-1b-it-sponsors"
+OUTPUT_DIR = "./t5gemma-s-s-ul2-it-sponsors"
 LORA_R = 16
 LORA_ALPHA = 32
 LORA_DROPOUT = 0.05
@@ -24,7 +26,7 @@ def main():
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    model = AutoModelForCausalLM.from_pretrained(
+    model = AutoModelForSeq2SeqLM.from_pretrained(
         MODEL_ID,
         device_map="auto",
         torch_dtype=torch.bfloat16,
@@ -44,9 +46,13 @@ def main():
         target_modules=LORA_TARGET_MODULES,
         lora_dropout=LORA_DROPOUT,
         bias="none",
-        task_type="CAUSAL_LM"
+        task_type="SEQ_2_SEQ_LM"
     )
 
+    print("Applying LoRA config to the model...")
+    model = get_peft_model(model, peft_config)
+    model.print_trainable_parameters()
+
     print("Model and LoRA config loaded. Setting up trainer...")
 
     training_args = SFTConfig(
@@ -57,7 +63,6 @@ def main():
         per_device_eval_batch_size=1,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=8,
-        gradient_checkpointing=True,
         learning_rate=1e-4,
         lr_scheduler_type="linear",
         optim="adamw_torch",
@@ -83,7 +88,6 @@ def main():
         model=model,
         processing_class=tokenizer,
         args=training_args,
-        peft_config=peft_config,
         train_dataset=train_dataset,
         eval_dataset=validation_dataset,
     )
diff --git a/inference.py b/inference.py
new file mode 100644
index 0000000..283fa1b
--- /dev/null
+++ b/inference.py
@@ -0,0 +1,53 @@
+import torch
+import re
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+MODEL_PATH = "./t5gemma-s-s-ul2-it-sponsors/final_checkpoint"
+BASE_MODEL_ID = "google/t5gemma-s-s-ul2-it"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+torch.set_float32_matmul_precision('high')
+
+print("Loading base model...")
+base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16)
+print("Loading LoRA weights...")
+model = PeftModel.from_pretrained(base_model, MODEL_PATH)
+
+print("Merging LoRA weights into the base model...")
+model = model.merge_and_unload()
+
+model.to(DEVICE)
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+def predict_timestamp(transcript_text, max_new_tokens=256):
+    input_text = (
+        "Instruction: Identify the timestamp range of the pre-recorded ad in the following transcript. "
+        "Output only the HH:MM:SS.mmm-HH:MM:SS.mmm range.\n\nTranscript:\n"
+        f"{transcript_text}"
+    )
+
+    inputs = tokenizer(input_text, return_tensors="pt").to(DEVICE)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    input_length = inputs.input_ids.shape[1]
+    newly_generated_ids = output_ids[0][input_length:]
+    output_text = tokenizer.decode(newly_generated_ids, skip_special_tokens=True)
+    return output_text
+
+if __name__ == "__main__":
+    # Single transcript
+    transcript = "[00:18:00.079 --> 00:18:01.119] Points on my bracket.\n[00:18:01.200 --> 00:18:02.640] That's the max score I can get.\n[00:18:02.799 --> 00:18:12.720] His max bracket score right now is 62 points, which is barely enough to compete with the people currently in the lead of our bracket challenge.\n[00:18:12.960 --> 00:18:15.039] So you're saying there's a chance?\n[00:18:15.359 --> 00:18:16.400] That's not what he's saying.\n[00:18:16.480 --> 00:18:17.759] That's not what he's saying at all.\n[00:18:17.920 --> 00:18:19.839] We rewind, we cut to him like actually.\n[00:18:20.160 --> 00:18:29.359] Actually, like mad scientist going through each team's roster, like seeing their potential draft status, and he's like trying to create and he's like, This is gonna be the best bracket we've ever seen.\n[00:18:29.599 --> 00:18:32.799] And he never put more effort to anything in his life, and this is the result.\n[00:18:33.039 --> 00:18:36.160] Well, it doesn't bode well for him that who did he pick to win it?\n[00:18:36.319 --> 00:18:38.319] He picked Oh, he picked the Zags.\n[00:18:38.400 --> 00:18:39.680] He picked on Zaga to win it, which\n[00:18:40.319 --> 00:18:43.839] Is not a completely insane pick, but they were an eight seed.\n[00:18:43.920 --> 00:18:45.440] They're a good tournament team.\n[00:18:45.519 --> 00:18:49.119] Like they're traditionally they're a good tournament team, but what was their seed this year?\n[00:18:49.279 --> 00:18:50.319] They were an eight seed.\n[00:18:50.480 --> 00:18:52.640] It's a bold strategy, but there's a chance.\n[00:18:53.359 --> 00:18:54.640] Mark, how's your bracket?\n[00:18:54.960 --> 00:18:55.680] Oh man.\n[00:18:55.839 --> 00:18:57.440] If I did one this year.\n[00:18:57.599 --> 00:18:59.680] Well, I've done one for the past few years.\n[00:19:00.240 --> 00:19:05.119] And I've done it completely by random, and each time I've beaten Tyler's bracket.\n[00:19:06.960 --> 00:19:09.119] Like not lying multiple times.\n[00:19:09.279 --> 00:19:10.559] I can't remember if it's every time.\n[00:19:10.799 --> 00:19:12.319] I'll bet that goes over well.\n[00:19:12.559 --> 00:19:19.759] But yeah, I just purely by random chance and I I don't just win, I crush, I crush him in points.\n[00:19:20.240 --> 00:19:22.240] Just absolutely destroy.\n[00:19:22.480 --> 00:19:27.359] I'm gonna put down better at sports than Tyler for a point for Mark.\n[00:19:27.599 --> 00:19:29.119] Yep, that's that's true actually.\n[00:19:29.279 --> 00:19:31.440] I'm surprised you guys didn't do one for Go.\n[00:19:31.680 --> 00:19:32.720] He's gone.\n[00:19:34.880 --> 00:19:35.920] Where'd he go?\n[00:19:36.160 --> 00:19:36.880] He's gone.\n[00:19:37.200 --> 00:19:38.000] Honorable.\n[00:19:38.160 --> 00:19:39.920] How many times have I started a rumor that?\n[00:19:40.240 --> 00:19:41.839] Tyler's dead on his podcast.\n[00:19:42.799 --> 00:19:44.079] At least one now.\n[00:19:44.319 --> 00:19:47.440] I forget where this is all a delay because I forgot where he's going.\n[00:19:47.599 --> 00:19:50.240] He told me many times, but he's not here.\n[00:19:50.480 --> 00:19:52.799] I know where he was until the middle of the night last night.\n[00:19:53.039 --> 00:19:54.400] That's very specific.\n[00:19:54.559 --> 00:19:58.079] Uh Bird, the porn guy, if you guys remember the porn guy from the porn episode.\n[00:19:58.240 --> 00:19:59.920] Uh I found out that he had to take Tyler's.\n[00:20:00.240 --> 00:20:01.200] to the airport last night.\n[00:20:01.359 --> 00:20:04.880] That's weird because I know where Bert lives and it's not near where Tyler lives.\n[00:20:05.119 --> 00:20:05.759] It is not.\n[00:20:05.920 --> 00:20:07.680] Meaning that Tyler was in fact gone.\n[00:20:07.920 --> 00:20:10.400] Maybe that should be the the topic of this episode.\n[00:20:10.480 --> 00:20:12.720] We need to get to the bottom of this ASAP.\n[00:20:12.960 --> 00:20:15.839] If we find out where he is by the end of the episode, do we get a point?\n[00:20:16.079 --> 00:20:16.400] Yeah.\n[00:20:16.640 --> 00:20:19.920] You can have two points if you find out where he is by the end of the episode.\n[00:20:20.480 --> 00:20:21.599] Everybody starts texting Tyler.\n[00:20:21.759 --> 00:20:22.640] Hey, where the fuck are you?\n[00:20:22.799 --> 00:20:27.279] I'm debating whether I text Tyler or do I reach for the Girl Scout cookies?\n[00:20:27.519 --> 00:20:29.839] Do I want the points or the sugar?\n[00:20:30.079 --> 00:20:32.400] Alright, let's say we can't just text him directly.\n[00:20:32.559 --> 00:20:34.160] That would be cheating.\n[00:20:34.480 --> 00:20:36.640] This episode is brought to you by Amazon.\n[00:20:36.799 --> 00:20:39.839] This off to college season, save on college, save the everyday.\n[00:20:40.160 --> 00:20:44.640] Literally every supply you need for school, like pens and stuff that you just assume you have.\n[00:20:44.799 --> 00:20:46.319] You know, I remember a story.\n[00:20:46.480 --> 00:20:49.519] My roommate and I decided to bunk our beds.\n[00:20:49.680 --> 00:20:53.119] So we just cut up some plastic hangers and just jammed them in there.\n[00:20:53.279 --> 00:20:56.160] We could have really used some metal pins.\n[00:20:56.400 --> 00:20:58.000] Wonder who that roommate could have been.\n[00:20:58.079 --> 00:20:59.920] So remember with Amazon's low off.\n[00:21:00.079 --> 00:21:02.559] Off to college prices, save on college, save the everyday.\n[00:21:02.720 --> 00:21:04.559] Shop off to college at Amazon.\n[00:21:04.720 --> 00:21:06.880] This episode is brought to you by Mentos Gum.\n[00:21:07.039 --> 00:21:09.119] Keep things fresh, it's important, right?\n[00:21:09.279 --> 00:21:10.799] And I'm not just talking about fresh breath.\n[00:21:10.880 --> 00:21:13.599] It's important to switch up your routine whenever you can.\n[00:21:13.759 --> 00:21:16.480] I just uh I'm the person who can't help but chew.\n[00:21:16.640 --> 00:21:18.640] You put up an event in your mouth, you're supposed to suck on it.\n[00:21:18.720 --> 00:21:18.880] I'm like\n[00:21:20.160 --> 00:21:21.759] Swallow, so I kinda need gum.\n[00:21:22.000 --> 00:21:23.759] You turn into a cartoon dog.\n[00:21:23.839 --> 00:21:24.240] I'm sorry.\n[00:21:25.519 --> 00:21:28.319] Next time we hang out, I'm giving you a mint just to see what happens.\n[00:21:28.480 --> 00:21:34.799] And of course, another way to refresh every day is with Mento's gum, available in a range of fresh flavors like spearmint, fresh mint, and strawberry.\n[00:21:35.039 --> 00:21:39.200] Mentos gum, yes, to fresh.\n[00:21:40.000 --> 00:21:42.400] Uh, should we get into the topic for today's episode?\n[00:21:42.640 --> 00:21:50.880] I gotta be honest, I looked and I'm only mediumly sure we haven't done something that's fairly similar to this, but I just sort of want to talk about it.\n[00:21:51.440 --> 00:21:59.359] I'm calling this episode probably not maybe something like uh Distractable Travel Guide, Cincinnati."
+    print(transcript)
+    predicted_range = predict_timestamp(transcript)
+    print("Predicted ad range:", predicted_range)
diff --git a/main.py b/main.py
index a58e02f..6887ac3 100644
--- a/main.py
+++ b/main.py
@@ -401,7 +401,7 @@ def process_transcript_json(input_json_path: str, output_jsonl_path: str):
         
         final_results.append({
             "id": prompt_id,
-            "text": full_prompt_for_finetuning,
+            "input": full_prompt_for_finetuning,
             "target": response_text.strip()
         })
 
diff --git a/reformat.py b/reformat.py
new file mode 100644
index 0000000..c21056c
--- /dev/null
+++ b/reformat.py
@@ -0,0 +1,66 @@
+import sys
+import json
+from transformers import AutoTokenizer
+from tqdm import tqdm
+
+# IMPORTANT: The tokenizer MUST match the student model you are fine-tuning.
+STUDENT_MODEL_ID = "google/t5gemma-s-s-ul2-it"
+
+def reformat_with_template(input_filepath: str, output_filepath: str):
+    """
+    Reads a JSONL file with 'text' (prompt) and 'target' (completion) fields
+    and reformats it into a single 'text' field using the model's chat template.
+    """
+    print(f"Loading tokenizer for {STUDENT_MODEL_ID}...")
+    tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_ID)
+    
+    print(f"Reformatting {input_filepath}...")
+    
+    with open(input_filepath, 'r', encoding='utf-8') as infile, \
+         open(output_filepath, 'w', encoding='utf-8') as outfile:
+        
+        # Count lines for tqdm progress bar
+        num_lines = sum(1 for line in open(input_filepath, 'r', encoding='utf-8'))
+        infile.seek(0) # Reset file pointer
+        
+        for line in tqdm(infile, total=num_lines, desc="Reformatting"):
+            original_entry = json.loads(line)
+            
+            # The original prompt is in the 'text' field
+            prompt_text = original_entry.get("text", "")
+            # The verified completion is in the 'target' field
+            completion_text = original_entry.get("target", "")
+            
+            # The prompt needs to end with the cue for the model to respond to
+            if "Timestamp:" not in prompt_text:
+                 prompt_text += "\n\nTimestamp:\n"
+
+            # Create the message structure that the chat template expects
+            messages = [
+                {"role": "user", "content": prompt_text},
+                {"role": "assistant", "content": completion_text}
+            ]
+            
+            # Apply the template to create the new, single text field
+            # add_generation_prompt=False is important for training data
+            formatted_text = tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False, 
+                add_generation_prompt=False
+            )
+            
+            new_entry = {
+                "id": original_entry.get("id"),
+                "text": formatted_text,
+                "target": completion_text # Keep target for reference if needed
+            }
+            outfile.write(json.dumps(new_entry) + "\n")
+
+    print(f"\nSuccessfully reformatted {num_lines} entries. Output saved to {output_filepath}.")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        print("Usage: python reformat_with_template.py <input_dataset.jsonl> <output_dataset_templated.jsonl>")
+        sys.exit(1)
+        
+    reformat_with_template(sys.argv[1], sys.argv[2])
diff --git a/split_dataset.py b/split_dataset.py
index 9ac7bbf..a161cc1 100644
--- a/split_dataset.py
+++ b/split_dataset.py
@@ -7,10 +7,12 @@ def split_dataset(
     input_filepath: str, 
     train_filepath: str = "train.jsonl", 
     validation_filepath: str = "validation.jsonl", 
-    split_ratio: float = 0.95
+    split_ratio: float = 0.95,
+    balance_ratio: float = 1.0
 ):
     """
-    Reads a .jsonl file, shuffles it, and splits it into training and validation files.
+    Reads a .jsonl file, balances the positive/negative examples, shuffles,
+    and splits it into training and validation files.
     """
     print(f"Loading data from {input_filepath}...")
     try:
@@ -20,23 +22,51 @@ def split_dataset(
         print(f"Error: Input file not found at {input_filepath}", file=sys.stderr)
         sys.exit(1)
 
-    # Shuffle the dataset to ensure random distribution
-    print("Shuffling data...")
-    random.shuffle(lines)
+    # 1. Separate into positive and negative examples
+    positive_examples = []
+    negative_examples = []
+    for line in lines:
+        try:
+            entry = json.loads(line)
+            if entry.get("target", "").strip():
+                positive_examples.append(line)
+            else:
+                negative_examples.append(line)
+        except json.JSONDecodeError:
+            print(f"Warning: Skipping malformed JSON line: {line.strip()}", file=sys.stderr)
 
-    # Determine the split point
-    split_index = int(len(lines) * split_ratio)
+    print(f"Found {len(positive_examples)} positive examples and {len(negative_examples)} negative examples.")
 
-    # Split the data
-    train_lines = lines[:split_index]
-    validation_lines = lines[split_index:]
+    # 2. Shuffle both lists independently
+    random.shuffle(positive_examples)
+    random.shuffle(negative_examples)
+
+    # 3. Balance the dataset based on the ratio
+    num_positives = len(positive_examples)
+    num_negatives_to_keep = int(num_positives * balance_ratio)
+
+    if len(negative_examples) < num_negatives_to_keep:
+        print(f"Warning: Not enough negative examples to meet a {balance_ratio}:1 ratio. "
+              f"Using all {len(negative_examples)} negative examples.", file=sys.stderr)
+        num_negatives_to_keep = len(negative_examples)
+
+    print(f"Balancing dataset with {num_positives} positive and {num_negatives_to_keep} negative examples.")
+    
+    balanced_lines = positive_examples + negative_examples[:num_negatives_to_keep]
+    
+    # 4. Final shuffle to mix positive and negative examples
+    print("Shuffling final balanced dataset...")
+    random.shuffle(balanced_lines)
+
+    # 5. Determine the split point and write files
+    split_index = int(len(balanced_lines) * split_ratio)
+    train_lines = balanced_lines[:split_index]
+    validation_lines = balanced_lines[split_index:]
 
-    # Write the training file
     print(f"Writing {len(train_lines)} lines to {train_filepath}...")
     with open(train_filepath, 'w', encoding='utf-8') as f:
         f.writelines(train_lines)
 
-    # Write the validation file
     print(f"Writing {len(validation_lines)} lines to {validation_filepath}...")
     with open(validation_filepath, 'w', encoding='utf-8') as f:
         f.writelines(validation_lines)
@@ -44,7 +74,9 @@ def split_dataset(
     print("\nSplitting complete.")
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Split a .jsonl dataset into training and validation sets.")
+    parser = argparse.ArgumentParser(
+        description="Balance and split a .jsonl dataset into training and validation sets."
+    )
     parser.add_argument("input_file", type=str, help="Path to the input .jsonl file.")
     parser.add_argument(
         "--split", 
@@ -52,6 +84,12 @@ if __name__ == "__main__":
         default=0.95, 
         help="Ratio for the training set split (e.g., 0.95 for a 95/5 split). Default is 0.95."
     )
+    parser.add_argument(
+        "--ratio",
+        type=float,
+        default=1.0,
+        help="Ratio of negative to positive examples (e.g., 1.0 for 1:1, 2.0 for 2:1). Default is 1.0."
+    )
     
     args = parser.parse_args()
     
@@ -59,4 +97,4 @@ if __name__ == "__main__":
         print("Error: Split ratio must be between 0 and 1.", file=sys.stderr)
         sys.exit(1)
         
-    split_dataset(args.input_file, split_ratio=args.split)
\ No newline at end of file
+    split_dataset(args.input_file, split_ratio=args.split, balance_ratio=args.ratio)
\ No newline at end of file