import spaces  # MUST come before torch / transformers
import torch
import gradio as gr

from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "superwhisper/s1-mini"

SYSTEM_PROMPT = (
    "You are a text normalizer for speech-to-text transcripts. "
    "The input begins with a control line specifying the styling, structure, "
    "and context settings; clean the transcript to match those settings "
    "and output only the cleaned text."
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()


@spaces.GPU(duration=30)
def clean_transcript(
    transcript: str,
    styling: str = "semi-formal",
    structure: str = "prose",
    context: str = "general",
) -> str:
    """Clean and normalize a raw ASR transcript.

    Applies punctuation, truecasing, filler removal, and formatting based on
    the selected styling, structure, and context settings.

    Args:
        transcript: Raw ASR transcript text (disfluent, unpunctuated).
        styling: Register / formality level (casual, semi-casual, semi-formal, formal).
        structure: Output structure (prose or lists).
        context: Context mode (general or email).

    Returns:
        Cleaned, normalized transcript as plain text.
    """
    control = f"[Styling: {styling}] [Structure: {structure}] [Context: {context}]"
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"{control}\n{transcript}"},
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    input_len = inputs.input_ids.shape[1]
    max_new = min(1024, int(input_len * 1.3) + 32)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new,
            do_sample=False,
        )
    generated = out[0][input_len:]
    result = tokenizer.decode(generated, skip_special_tokens=True).strip()
    return result


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# S1-mini · ASR Transcript Cleanup\n"
        "A 0.6B Qwen3-based model that cleans raw speech-to-text transcripts: "
        "removes fillers, resolves self-corrections, adds punctuation and truecasing, "
        "and formats numbers/dates/emails — all controlled by styling and structure settings.\n\n"
        "Model: [superwhisper/s1-mini](https://huggingface.co/superwhisper/s1-mini)"
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            transcript_input = gr.Textbox(
                label="Raw ASR Transcript",
                placeholder="Paste raw, unpunctuated speech-to-text output here…",
                lines=6,
                scale=4,
            )
        with gr.Row():
            styling = gr.Dropdown(
                choices=["casual", "semi-casual", "semi-formal", "formal"],
                value="semi-formal",
                label="Styling",
                scale=1,
            )
            structure = gr.Dropdown(
                choices=["prose", "lists"],
                value="prose",
                label="Structure",
                scale=1,
            )
            context = gr.Dropdown(
                choices=["general", "email"],
                value="general",
                label="Context",
                scale=1,
            )
        run_btn = gr.Button("Clean Transcript", variant="primary")
        output = gr.Textbox(
            label="Cleaned Transcript",
            lines=6,
        )

    run_btn.click(
        fn=clean_transcript,
        inputs=[transcript_input, styling, structure, context],
        outputs=output,
        api_name="clean",
    )

    gr.Examples(
        examples=[
            ["so um i need to like send the the report by uh friday no wait make that thursday"],
            ["hey can you like um check the the numbers for q3 and also um make sure the the spreadsheet is up to date"],
            ["hi john um i was wondering if you could um send me the the quarterly report by end of day friday thanks"],
            ["so the meeting is at um three pm on tuesday and we need to like bring the the slides and also the budget numbers"],
        ],
        inputs=transcript_input,
        outputs=output,
        fn=clean_transcript,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)