Create app.py

app.py

@@ -0,0 +1,199 @@
+import os
+import json
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+MODEL_ID = os.getenv("MODEL_ID", "oleh13/ord-retro-qwen25-15b-merged-fp16")
+
+hf_token = os.getenv("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
+
+app = FastAPI(
+    title="ORD Retrosynthesis Model API",
+    version="1.0.0",
+)
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+
+class GenerateRequest(BaseModel):
+    prompt: Optional[str] = None
+    messages: Optional[List[ChatMessage]] = None
+
+    max_new_tokens: int = Field(default=1200, ge=1, le=4096)
+    temperature: float = Field(default=0.1, ge=0.0, le=2.0)
+    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
+    repetition_penalty: float = Field(default=1.05, ge=0.5, le=2.0)
+    do_sample: bool = True
+
+    return_json_only: bool = True
+
+
+class GenerateResponse(BaseModel):
+    text: str
+    parsed_json: Optional[Union[Dict[str, Any], List[Any]]] = None
+    raw_output: str
+    model_id: str
+
+
+def extract_json_object(text: str) -> str:
+    text = text.strip()
+
+    first_obj = text.find("{")
+    last_obj = text.rfind("}")
+
+    first_arr = text.find("[")
+    last_arr = text.rfind("]")
+
+    obj_valid = first_obj != -1 and last_obj != -1 and last_obj > first_obj
+    arr_valid = first_arr != -1 and last_arr != -1 and last_arr > first_arr
+
+    if obj_valid and arr_valid:
+        if first_obj < first_arr:
+            return text[first_obj:last_obj + 1]
+        return text[first_arr:last_arr + 1]
+
+    if obj_valid:
+        return text[first_obj:last_obj + 1]
+
+    if arr_valid:
+        return text[first_arr:last_arr + 1]
+
+    raise ValueError("No JSON object or array found in model output.")
+
+
+print(f"Loading tokenizer: {MODEL_ID}")
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+print(f"Loading model: {MODEL_ID}")
+
+if torch.cuda.is_available():
+    torch_dtype = torch.bfloat16
+    device_map = "auto"
+else:
+    torch_dtype = torch.float32
+    device_map = "cpu"
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch_dtype,
+    device_map=device_map,
+    trust_remote_code=True,
+    low_cpu_mem_usage=True,
+)
+
+model.eval()
+
+print("Model loaded.")
+print("CUDA:", torch.cuda.is_available())
+
+
+@app.get("/")
+def root():
+    return {
+        "status": "ok",
+        "model_id": MODEL_ID,
+        "cuda": torch.cuda.is_available(),
+    }
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "model_id": MODEL_ID,
+        "cuda": torch.cuda.is_available(),
+    }
+
+
+@app.post("/generate", response_model=GenerateResponse)
+def generate(req: GenerateRequest):
+    if req.messages and req.prompt:
+        raise HTTPException(
+            status_code=400,
+            detail="Send either 'prompt' or 'messages', not both.",
+        )
+
+    if not req.messages and not req.prompt:
+        raise HTTPException(
+            status_code=400,
+            detail="Send either 'prompt' or 'messages'.",
+        )
+
+    if req.messages:
+        messages = [m.model_dump() for m in req.messages]
+
+        prompt_text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+    else:
+        prompt_text = req.prompt
+
+    inputs = tokenizer(
+        [prompt_text],
+        return_tensors="pt",
+    )
+
+    if torch.cuda.is_available():
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=req.max_new_tokens,
+            temperature=req.temperature,
+            top_p=req.top_p,
+            repetition_penalty=req.repetition_penalty,
+            do_sample=req.do_sample,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+
+    raw_output = tokenizer.decode(
+        outputs[0][inputs["input_ids"].shape[-1]:],
+        skip_special_tokens=True,
+    ).strip()
+
+    parsed_json = None
+    final_text = raw_output
+
+    if req.return_json_only:
+        try:
+            json_text = extract_json_object(raw_output)
+            parsed_json = json.loads(json_text)
+            final_text = json.dumps(parsed_json, ensure_ascii=False, indent=2)
+        except Exception as exc:
+            raise HTTPException(
+                status_code=422,
+                detail={
+                    "message": "Model did not return valid JSON.",
+                    "error": str(exc),
+                    "raw_output": raw_output,
+                },
+            )
+
+    return GenerateResponse(
+        text=final_text,
+        parsed_json=parsed_json,
+        raw_output=raw_output,
+        model_id=MODEL_ID,
+    )