Upload app.py with huggingface_hub

app.py

@@ -1,19 +1,36 @@
 import html
+import json
+import os
 import time
+import uuid
 
+os.environ.setdefault("GRADIO_SSR_MODE", "false")
+
+import fastapi
 import spaces
 import torch
 import gradio as gr
+from fastapi.responses import JSONResponse, StreamingResponse
+from gradio import Server
+from gradio.context import LocalContext
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+try:
+    from gradio.route_utils import Request as GradioRequest
+except ImportError:
+    GradioRequest = None
+
 from pasita import infer_kind, convert_to_markdown
 
 MODEL_ID = "OpceanAI/PASITA"
 MODEL_CARD = "https://huggingface.co/OpceanAI/PASITA"
+MAX_INPUT_CHARS = 60_000
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 tokenizer.padding_side = "left"
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype="auto").to("cuda")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype="auto").to(DEVICE)
 model.eval()
 
 EMPTY_MD = "_The converted Markdown appears here. Paste a document on the left and press Convert._"
@@ -48,6 +65,25 @@ def _estimate_duration(text, num_candidates=3, structured_fallback=True, *args,
     return min(30, max(12, 8 + len(text or "") // 200 + 4 * (int(num_candidates) - 1)))
 
 
+def _run_harness(text: str, num_candidates: int, structured_fallback: bool) -> dict:
+    started = time.perf_counter()
+    result = convert_to_markdown(
+        model,
+        tokenizer,
+        text,
+        num_candidates=max(1, min(3, int(num_candidates))),
+        do_sample=True,
+        temperature=0.3,
+        top_p=0.95,
+        seed=0,
+        use_cache=True,
+        structured_fallback=bool(structured_fallback),
+        skip_model_for_structured=True,
+    )
+    result["elapsed"] = time.perf_counter() - started
+    return result
+
+
 @spaces.GPU(duration=_estimate_duration)
 def convert(text: str, num_candidates: int = 3, structured_fallback: bool = True) -> tuple[str, str, str]:
     """Convert a plain-text document into faithful GitHub-Flavored Markdown.
@@ -68,27 +104,14 @@ def convert(text: str, num_candidates: int = 3, structured_fallback: bool = True
     if not text:
         return EMPTY_MD, "", _status("READY / PASTE A DOCUMENT TO CONVERT")
 
-    started = time.perf_counter()
     try:
-        result = convert_to_markdown(
-            model,
-            tokenizer,
-            text,
-            num_candidates=max(1, min(3, int(num_candidates))),
-            do_sample=True,
-            temperature=0.3,
-            top_p=0.95,
-            seed=0,
-            use_cache=True,
-            structured_fallback=bool(structured_fallback),
-            skip_model_for_structured=True,
-        )
+        result = _run_harness(text, num_candidates, structured_fallback)
     except Exception as exc:
         message = html.escape(f"{type(exc).__name__}: {exc}")
         return (
             f"> Conversion failed.\n>\n> `{message}`",
             "",
-            _status(f"ERROR / {type(exc).__name__}", "error", f"{time.perf_counter() - started:.1f}S"),
+            _status(f"ERROR / {type(exc).__name__}", "error"),
         )
 
     metrics = result["metrics"]
@@ -100,15 +123,243 @@ def convert(text: str, num_candidates: int = 3, structured_fallback: bool = True
         "passthrough": "SOURCE PASSTHROUGH",
     }
     source_label = routes.get(result.get("route"), "PASITA + RERANK")
-    taken = result["seconds"] if result["seconds"] else time.perf_counter() - started
+    taken = result["seconds"] if result["seconds"] else result["elapsed"]
     detail = f"{source_label} / {numbers} / {overlap} / {taken:.1f}S"
     return result["markdown"], result["markdown"], _status("DONE", "ready", detail)
 
 
+@spaces.GPU(duration=_estimate_duration)
+def gpu_convert(text: str, num_candidates: int = 3, structured_fallback: bool = True) -> dict:
+    """Convert a document and return the full conversion result.
+
+    Same pipeline as the UI conversion. Used by the /v1 OpenAI-style API routes.
+
+    Args:
+        text: The document to convert.
+        num_candidates: How many PASITA samples to generate and rerank, from 1 to 3.
+        structured_fallback: Use the deterministic converter for HTML and delimited data.
+    Returns:
+        The full result dict: markdown, route, metrics, timing and candidate info.
+    """
+    return _run_harness(text, num_candidates, structured_fallback)
+
+
+COUNTER_JS = """(text) => {
+  const t = (text || '').trim();
+  const words = t ? t.split(/\\s+/).length : 0;
+  const chars = (text || '').length;
+  return '<span class="p-label">PLAIN TEXT</span><span class="p-count">' + chars.toLocaleString() + ' CHARS / ' + words.toLocaleString() + ' WORDS</span>';
+}"""
+
+API_DESCRIPTION = """
+OpenAI-style REST API for [PASITA](https://huggingface.co/OpceanAI/PASITA), an 88M
+from-scratch model that turns plain text into faithful Markdown. No API key is required.
+Works with the OpenAI SDK by pointing `base_url` at `/v1`. The Gradio API and the MCP
+server stay available at `?view=api`.
+""".strip()
+
+
+def _api_error(message: str, error_type: str = "invalid_request_error", status: int = 400, code: str | None = None) -> JSONResponse:
+    return JSONResponse(
+        status_code=status,
+        content={"error": {"message": message, "type": error_type, "param": None, "code": code}},
+    )
+
+
+def _with_request_context(request: fastapi.Request, text: str, candidates: int, fallback: bool) -> dict:
+    if GradioRequest is not None:
+        gr_request = GradioRequest(request=request, session_hash="api")
+        token = LocalContext.request.set(gr_request)
+        try:
+            return gpu_convert(text, candidates, fallback)
+        finally:
+            LocalContext.request.reset(token)
+    return gpu_convert(text, candidates, fallback)
+
+
+def _token_counts(text: str, markdown: str) -> tuple[int, int]:
+    prompt_tokens = len(tokenizer(text, add_special_tokens=True)["input_ids"])
+    completion_tokens = len(tokenizer(markdown, add_special_tokens=False)["input_ids"])
+    return prompt_tokens, completion_tokens
+
+
+class ChatMessage(BaseModel):
+    role: str = "user"
+    content: str | list | None = None
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str | None = "pasita-v1"
+    messages: list[ChatMessage]
+    stream: bool = False
+    candidates: int = 3
+    structured_fallback: bool = True
+
+
+class ConvertRequest(BaseModel):
+    text: str
+    candidates: int = 3
+    structured_fallback: bool = True
+
+
+server = Server(title="PASITA API", version="1.0.0", description=API_DESCRIPTION)
+
+
+@server.exception_handler(fastapi.exceptions.RequestValidationError)
+async def validation_error_handler(request: fastapi.Request, exc: fastapi.exceptions.RequestValidationError):
+    issues = "; ".join(f"{'.'.join(str(l) for l in e.get('loc', []))}: {e.get('msg', '')}" for e in exc.errors()[:3])
+    return _api_error(f"Invalid request body: {issues}", status=400)
+
+
+@server.get("/v1/models")
+def list_models() -> dict:
+    return {
+        "object": "list",
+        "data": [{"id": "pasita-v1", "object": "model", "created": 1760000000, "owned_by": "OpceanAI"}],
+    }
+
+
+@server.get("/v1/health")
+def health() -> dict:
+    return {"status": "ok", "model": "pasita-v1"}
+
+
+def _validate(body: ChatCompletionRequest | ConvertRequest, text: str | None) -> JSONResponse | None:
+    if text is None or not text.strip():
+        return _api_error("No input text was provided.", code="empty_input")
+    if len(text) > MAX_INPUT_CHARS:
+        return _api_error(
+            f"The document is too long ({len(text):,} chars). The limit is {MAX_INPUT_CHARS:,} characters.",
+            code="input_too_long",
+        )
+    candidates = getattr(body, "candidates", 3)
+    if not 1 <= int(candidates) <= 3:
+        return _api_error("candidates must be between 1 and 3.", code="invalid_candidates")
+    return None
+
+
+def _result_payload(result: dict, text: str, started: float) -> dict:
+    markdown = result["markdown"]
+    prompt_tokens, completion_tokens = _token_counts(text, markdown)
+    return {
+        "markdown": markdown,
+        "route": result.get("route", "model"),
+        "metrics": result["metrics"],
+        "candidates": len(result.get("candidates", [])),
+        "seconds": round(result["seconds"] or result["elapsed"], 2),
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+        "elapsed": round(time.perf_counter() - started, 2),
+    }
+
+
+@server.post("/v1/convert")
+def v1_convert(body: ConvertRequest, request: fastapi.Request) -> dict | JSONResponse:
+    error = _validate(body, body.text)
+    if error is not None:
+        return error
+    started = time.perf_counter()
+    try:
+        result = _with_request_context(request, body.text.strip(), int(body.candidates), bool(body.structured_fallback))
+    except Exception as exc:
+        return _api_error(f"{type(exc).__name__}: {exc}", "server_error", status=500)
+    payload = _result_payload(result, body.text.strip(), started)
+    payload.update({"id": f"conv-{uuid.uuid4().hex[:24]}", "object": "pasita.conversion", "created": int(time.time()), "model": "pasita-v1"})
+    return payload
+
+
+def _extract_user_text(messages: list[ChatMessage]) -> str | None:
+    for message in reversed(messages):
+        if message.role != "user":
+            continue
+        content = message.content
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts = []
+            for part in content:
+                if isinstance(part, dict) and part.get("type") == "text":
+                    parts.append(str(part.get("text", "")))
+                elif isinstance(part, str):
+                    parts.append(part)
+            return "\n".join(p for p in parts if p)
+        return None
+    return None
+
+
 def reset() -> tuple[str, str, str, str]:
     return EMPTY_MD, "", _status("READY"), _counter("")
 
 
+@server.post("/v1/chat/completions")
+def chat_completions(body: ChatCompletionRequest, request: fastapi.Request) -> dict | JSONResponse | StreamingResponse:
+    text = _extract_user_text(body.messages)
+    error = _validate(body, text)
+    if error is not None:
+        return error
+
+    started = time.perf_counter()
+    try:
+        result = _with_request_context(request, text.strip(), int(body.candidates), bool(body.structured_fallback))
+    except Exception as exc:
+        return _api_error(f"{type(exc).__name__}: {exc}", "server_error", status=500)
+
+    payload = _result_payload(result, text.strip(), started)
+    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
+    created = int(time.time())
+    markdown = payload["markdown"]
+
+    if body.stream:
+        def sse():
+            def chunk(delta, finish=None, usage=None):
+                data = {
+                    "id": completion_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": "pasita-v1",
+                    "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
+                }
+                if usage is not None:
+                    data["usage"] = usage
+                return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
+
+            yield chunk({"role": "assistant", "content": ""})
+            yield chunk({"content": markdown})
+            yield chunk({}, finish="stop", usage=payload["usage"])
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(
+            sse(),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+        )
+
+    return {
+        "id": completion_id,
+        "object": "chat.completion",
+        "created": created,
+        "model": "pasita-v1",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": markdown},
+                "logprobs": None,
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": payload["usage"],
+        "pasita": {
+            "route": payload["route"],
+            "metrics": payload["metrics"],
+            "candidates": payload["candidates"],
+            "seconds": payload["seconds"],
+        },
+    }
+
+
 BUSY_JS = """(text, num_candidates, structured_fallback, md, source, status) => {
   const n = Number(num_candidates) || 1;
   const label = n > 1 ? 'PASITA ' + n + ' SAMPLES / RERANK' : 'PASITA SAMPLE / RERANK';
@@ -469,4 +720,4 @@ with gr.Blocks(title="PASITA: plain text to Markdown") as demo:
     clear.click(reset, inputs=None, outputs=[md_out, source, status, label_in], api_visibility="undocumented")
 
 demo.queue()
-demo.launch(theme=theme, css=CSS, mcp_server=True)
+demo.launch(_app=server, theme=theme, css=CSS, mcp_server=True)