"""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  # ZeroGPU hardware refuses to start without one decorated function

    @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:  # both models run on the CPU, one request at a time, so visitors need no GPU quota
        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()