| 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("/chat/completions", 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, |
| ) |
