from unsloth import FastLanguageModel, is_bfloat16_supported
MAX_SEQ_LEN = 2048
MAX_PROMPT_LEN = 512
MAX_COMPLETION_LEN = 512
model, tokenizer = FastLanguageModel.from_pretrained(
model_name="pankaj0/llama-3.2-3b-asr-refiner-merged",
max_seq_length=MAX_SEQ_LEN,
load_in_4bit=True,
dtype="float16",
)
from unsloth.chat_templates import get_chat_template
# Setup Llama-3.x chat formatting
tokenizer = get_chat_template(
tokenizer,
chat_template="llama-3.1",
)
SYSTEM_PROMPT = (
"You are an expert speech-to-text post-processor. "
"Fix transcription misspellings, add correct casing and punctuation, "
"and remove speech disfluencies (e.g., 'um', 'uh', false starts). "
"Preserve the speaker's original meaning and exact phrasing. Do not summarize or rewrite."
)
FastLanguageModel.for_inference(model)
# Define multiple complex, long, and noisy evaluation samples
evaluation_samples = [
{
"description": "Sample 1: Heavy disfluencies, repetitions, and grammar errors",
"raw": "so basically uh the the projt dedline is next mondy. rght? but like... we we haven't even got the feedback from the client yet and i'm like super worried we gonna miss it. um, can you check with john?"
},
{
"description": "Sample 2: Highly technical, microphone noise, lack of punctuation, and phonetic spelling",
"raw": "ok so the database connection string is like totally broken. bzz. we are getting connection timeout errors on port five four three two and i think the host name should be db dot internal dot local instead of localhost rght? pls check the yaml config immediately"
},
{
"description": "Sample 3: Long conversational rambling with false starts and slang",
"raw": "well i was thinking uh maybe we should... wait actually no. let's just go with the first option. because the second option is like way too expensive and honestly speaking we don't have the bandwidth to build a custom api from scratch right now anyway. so yeah let's schedule a meeting tomorrow at nine am to align the team."
}
]
# Loop through each sample and generate corrections
for idx, sample in enumerate(evaluation_samples, 1):
print(f"=== Test {idx}: {sample['description']} ===")
print(f"Raw Input:\n\"{sample['raw']}\"\n")
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": f"Raw transcript:\n{sample['raw']}"}
]
inputs = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt"
).to("cuda")
outputs = model.generate(
input_ids=inputs,
max_new_tokens=256,
temperature=0.1, # Low temperature for precise corrections
use_cache=True
)
# Decode output (skipping prompt tokens)
response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
print(f"Corrected:\n\"{response}\"\n")
print("-" * 80 + "\n")
output
=== Test 1: Sample 1: Heavy disfluencies, repetitions, and grammar errors ===
Raw Input:
"so basically uh the the projt dedline is next mondy. rght?
but like... we we haven't even got the feedback from the
client yet and i'm like super worried we gonna miss it.
um, can you check with john?"
Corrected:
"So basically, the project deadline is next Monday. Right?
But like... we haven't even got the feedback from the
client yet, and I'm like super worried we're gonna miss it.
Can you check with John?"
--------------------------------------------------------------------------------
=== Test 2: Sample 2: Highly technical, microphone noise, lack of punctuation, and phonetic spelling ===
Raw Input:
"ok so the database connection string is like totally broken.
bzz. we are getting connection timeout errors on port five
four three two and i think the host name should
be db dot internal dot local instead of localhost rght?
pls check the yaml config immediately"
Corrected:
"The database connection string is like totally broken.
We are getting connection timeout errors on
port five four three two and I think the host name
should be `db.internal.local` instead of `localhost`, right?
Please check the YAML config immediately."
--------------------------------------------------------------------------------
=== Test 3: Sample 3: Long conversational rambling with false starts and slang ===
Raw Input:
"well i was thinking uh maybe we should... wait actually no.
let's just go with the first option. because the second
option is like way too expensive and honestly speaking we
don't have the bandwidth to build a custom api from scratch
right now anyway. so yeah let's schedule a meeting tomorrow
at nine am to align the team."
Corrected:
"Well, I was thinking maybe we should... wait, actually,
no. Let's just go with the first option. Because the second
option is like way too expensive, and honestly
speaking, we don't have the bandwidth to build a
custom API from scratch right now, anyway.
So, yeah, let's schedule a meeting tomorrow at nine
am to align the team."
--------------------------------------------------------------------------------