Transcribe and Correct with AI

import os, warnings
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import torch

warnings.filterwarnings("ignore", message=".*deprecated.*")

# --- Settings ---
MODEL    = "KBLab/kb-whisper-large"
DEVICE   = "cuda:0" if torch.cuda.is_available() else "cpu"
DTYPE    = torch.float16 if torch.cuda.is_available() else torch.float32
DIR      = "chunks"
OUTFILE  = "transcript_with_timestamps.txt"
CHUNK_DURATION_S = 30

# Load the model and processor
print(f"Loading model {MODEL} to {DEVICE}...")
model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL, torch_dtype=DTYPE, use_safetensors=True).to(DEVICE)
processor = AutoProcessor.from_pretrained(MODEL)
asr_pipeline = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    device=DEVICE,
    torch_dtype=DTYPE
)

# Helper function to format time
def format_time(seconds):
    h, r = divmod(seconds, 3600)
    m, s = divmod(r, 60)
    return f"{int(h):02d}:{int(m):02d}:{int(s):02d}"

# Fetch and sort the audio files
files = sorted([f for f in os.listdir(DIR) if f.endswith(".wav")])

# Open a file to write the result to
with open(OUTFILE, "w", encoding="utf-8") as out:
    for filename in files:
        chunk_index = int(filename.split("_")[1].split(".")[0])
        start_time = chunk_index * CHUNK_DURATION_S
        end_time = start_time + CHUNK_DURATION_S

        filepath = os.path.join(DIR, filename)
        print(f"Transcribing {filename}...")

        result = asr_pipeline(filepath, generate_kwargs={"language": "sv"})
        text = result["text"].strip()

        out.write(f"[{format_time(start_time)} - {format_time(end_time)}]\n{text}\n\n")

print(f"Transcription complete! Result saved in: {OUTFILE}")

import re
from transformers import pipeline

# --- Settings ---
INPUT_FILE  = "transcript_with_timestamps.txt"
OUTPUT_FILE = "transcript_punctuated.txt"
MODEL = "sdadas/byt5-text-correction"
BATCH_SIZE = 8 # Process 8 text blocks at a time for efficiency

# Function to ensure capital letter after period
def capitalize_sentences(text):
    parts = re.split(r'([.?!]\s*)', text)
    return "".join(p.capitalize() for p in parts)

print(f"Loading punctuation model {MODEL}...")
punctuation_pipeline = pipeline(
    "text2text-generation",
    model=MODEL,
    tokenizer=MODEL,
    device=0, # Use first GPU
    batch_size=BATCH_SIZE
)

# Read in the raw text
with open(INPUT_FILE, 'r', encoding='utf-8') as f:
    lines = f.read().splitlines()

# Separate timestamps from text to be processed
timestamps, texts_to_process = [], []
for line in lines:
    if line.startswith("[") or not line.strip():
        timestamps.append(line)
    else:
        timestamps.append(None) # Marker for text
        texts_to_process.append(line)

print(f"Correcting {len(texts_to_process)} text segments...")
corrected_texts = punctuation_pipeline(texts_to_process)

# Assemble the final text
output_lines = []
text_index = 0
for ts in timestamps:
    if ts is not None:
        output_lines.append(ts)
    else:
        # Fetch corrected text and apply capitalization
        corrected_text = corrected_texts[text_index]['generated_text'].strip()
        final_text = capitalize_sentences(corrected_text)
        output_lines.append(final_text)
        text_index += 1

with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
    f.write("\n".join(output_lines))

print(f"Punctuation complete! Result saved in: {OUTPUT_FILE}")