Upload 8 files

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 8000
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

__init__.py

@@ -0,0 +1,3 @@
+from app.main import app
+
+__all__ = ["app"]

config.py

@@ -0,0 +1,2 @@
+MODEL_NAME = "mjpsm/qwen3-0.6-bash-experiment-model-final-merged"
+MAX_NEW_TOKENS = 128

main.py

@@ -0,0 +1,26 @@
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from app.model_loader import load_model
+from app.routes import router
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    tokenizer, model = load_model()
+    app.state.tokenizer = tokenizer
+    app.state.model = model
+
+    yield
+
+    app.state.tokenizer = None
+    app.state.model = None
+
+
+app = FastAPI(
+    title="Qwen Bash Tool Calling API",
+    lifespan=lifespan,
+)
+
+app.include_router(router)

model_loader.py

@@ -0,0 +1,23 @@
+from app.config import MODEL_NAME
+
+
+def load_model():
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        trust_remote_code=True,
+        extra_special_tokens={},
+    )
+
+    print("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+    model.eval()
+
+    print("Model loaded.")
+    return tokenizer, model

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+torch
+transformers
+accelerate
+sentencepiece
+safetensors

routes.py

@@ -0,0 +1,105 @@
+import re
+import time
+
+from fastapi import APIRouter, Request
+
+from app.config import MAX_NEW_TOKENS, MODEL_NAME
+from app.schemas import PredictionResponse, PromptRequest
+
+
+router = APIRouter()
+
+COMMAND_PATTERN = re.compile(
+    r'"command"\s*:\s*"([^"]+)"',
+)
+
+
+@router.get("/")
+def root():
+    return {
+        "status": "running",
+    }
+
+
+@router.get("/health")
+def health(request: Request):
+    model_loaded = (
+        hasattr(request.app.state, "model")
+        and hasattr(request.app.state, "tokenizer")
+        and request.app.state.model is not None
+        and request.app.state.tokenizer is not None
+    )
+
+    return {
+        "status": "healthy",
+        "model_loaded": model_loaded,
+        "model_name": MODEL_NAME,
+    }
+
+
+@router.get("/model-info")
+def model_info():
+    return {
+        "model_name": MODEL_NAME,
+    }
+
+
+@router.post("/predict", response_model=PredictionResponse)
+def predict(payload: PromptRequest, request: Request):
+    import torch
+
+    start_time = time.time()
+
+    tokenizer = request.app.state.tokenizer
+    model = request.app.state.model
+
+    messages = [
+        {
+            "role": "user",
+            "content": payload.prompt,
+        }
+    ]
+
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+    ).to(model.device)
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False,
+        )
+
+    prompt_token_count = inputs["input_ids"].shape[1]
+    generated_tokens = output[0][prompt_token_count:]
+
+    response = tokenizer.decode(
+        generated_tokens,
+        skip_special_tokens=True,
+    )
+
+    command = None
+    match = COMMAND_PATTERN.search(response)
+
+    if match:
+        command = match.group(1)
+
+    latency_seconds = round(
+        time.time() - start_time,
+        3,
+    )
+
+    return PredictionResponse(
+        prompt=payload.prompt,
+        command=command,
+        raw_output=response,
+        latency_seconds=latency_seconds,
+    )

schemas.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+
+
+class PromptRequest(BaseModel):
+    prompt: str
+
+
+class PredictionResponse(BaseModel):
+    prompt: str
+    command: str | None
+    raw_output: str
+    latency_seconds: float