| """WriteFix demo Space: both models, picked with a switch. The website's "WF-1 Large" toggle calls the `fix` API here.""" |
| import threading |
| import gradio as gr |
| from huggingface_hub import hf_hub_download |
| from infer import Fixer |
|
|
| try: |
| import spaces |
|
|
| @spaces.GPU(duration=5) |
| def _zerogpu_placeholder(): |
| return "ok" |
| except Exception: |
| pass |
|
|
| REPO = "rishi38/writefix-models" |
| MODELS = {"WF-1 Mini": "wf-1-mini.pt", "WF-1 Large": "wf-1-large.pt"} |
| MAX_CHARS = 1500 |
| fixers, lock = {}, threading.Lock() |
|
|
|
|
| def get(name): |
| with lock: |
| if name not in fixers: |
| fixers[name] = Fixer(hf_hub_download(REPO, MODELS[name]), device="cpu") |
| return fixers[name] |
|
|
|
|
| def fix(text, model="WF-1 Large"): |
| if model not in MODELS: |
| model = "WF-1 Large" |
| f = get(model) |
| with lock: |
| return f.fix(text[:MAX_CHARS]) |
|
|
|
|
| for name in MODELS: |
| get(name) |
|
|
| demo = gr.Interface( |
| fn=fix, |
| inputs=[gr.Textbox(lines=8, label="Messy text", placeholder="type like you normally do no fullstops no commas spelling mistaks and all"), |
| gr.Radio(list(MODELS), value="WF-1 Large", label="Model", info="Mini: 10.5M parameters, 21 MB. Large: 94M parameters, 187 MB.")], |
| outputs=gr.Textbox(lines=8, label="Fixed"), |
| title="WriteFix", |
| description="Two correction models trained from scratch, with no pretrained weights. They fix spelling, punctuation, capitals and common grammar slips. English only.", |
| examples=[["i dont know wat happend but the server stoped responding after i deploy it", "WF-1 Large"], |
| ["lets meet tommorow at 5 and discus the plan for hackathon", "WF-1 Mini"]], |
| cache_examples=False, |
| api_name="fix", |
| ) |
| demo.queue(default_concurrency_limit=2).launch() |
|
|
