Hugging Face's logo

Spaces:
superwhisper
/
s1-mini-transcript-cleanup
Running on Zero

App Files Community
s1-mini-transcript-cleanup
File size: 4,768 Bytes
af99ebb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a45bdbf
af99ebb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a45bdbf
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
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)