try finetuning encoder-decoder model like t5gemmaHEADmaster

Unfortunately once it finishes the finetuning process, it doesn't
actually output anything. With decoder-only models like small gemma3, it
also doesn't output anything meaningful. I'm tempted to just classify
each token as an ad or not, finetune an encoder like some long-context
BERT, and then a post-processing script to convert back to tokens and
their timestamps. May end up being more performant and efficient too...

1 files changed, 53 insertions, 0 deletions

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)