Spaces:
Running on Zero
Running on Zero
File size: 4,018 Bytes
cbee686 fd88749 cbee686 fd88749 cbee686 fd88749 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import base64
import json
import re
import time
import boto3
# Some Gemma checkpoints wrap an internal reasoning/scratchpad span in these
# special-vocab tokens before the real answer; decoded as literal text here
# instead of being suppressed, so it must be stripped before display.
_THINKING_RE = re.compile(r"<unused94>.*?<unused95>", re.DOTALL)
REGION = "us-east-1"
BUCKET = "medgemma-async-2026"
ENDPOINT = "medgemma-vllm-endpoint"
_sagemaker = boto3.client("sagemaker-runtime", region_name=REGION)
_s3 = boto3.client("s3", region_name=REGION)
def encode_image_b64(image_bytes: bytes) -> str:
return base64.b64encode(image_bytes).decode()
def build_payload(messages: list[dict], image_b64: str | None,
max_new_tokens: int = 300, temperature: float = 0.2) -> dict:
"""DJL-LMI's chat-completions path (djl_python/chat_completions) is triggered by a
top-level "messages" key and expects OpenAI-style params (max_tokens, temperature)
at the top level too -- NOT nested under "parameters", and NOT under an "inputs" key
(that key is reserved for the plain text-generation handler, which rejects lists).
"""
if image_b64:
last = messages[-1]
last["content"] = [
{"type": "text", "text": last["content"]},
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
]
return {
"messages": messages,
"max_tokens": max_new_tokens,
"temperature": temperature,
}
class InferenceTimeout(Exception):
pass
def invoke_and_wait(payload: dict, timeout_s: int = 1200, poll_interval_s: int = 5) -> dict:
"""Submit an async inference job and poll S3 for the result.
Cold start (scale-to-zero -> instance launch -> model load) can take
10-15+ minutes on ml.g5.xlarge; timeout_s defaults high to cover that.
"""
input_key = f"input/chat_{int(time.time() * 1000)}.json"
_s3.put_object(Bucket=BUCKET, Key=input_key, Body=json.dumps(payload).encode())
response = _sagemaker.invoke_endpoint_async(
EndpointName=ENDPOINT,
InputLocation=f"s3://{BUCKET}/{input_key}",
ContentType="application/json",
)
output_key = response["OutputLocation"].replace(f"s3://{BUCKET}/", "")
failure_key = output_key.replace("output/", "failures/")
deadline = time.time() + timeout_s
while time.time() < deadline:
try:
obj = _s3.get_object(Bucket=BUCKET, Key=output_key)
return json.loads(obj["Body"].read())
except _s3.exceptions.NoSuchKey:
pass
try:
obj = _s3.get_object(Bucket=BUCKET, Key=failure_key)
raise RuntimeError(f"Inference failed: {obj['Body'].read().decode()}")
except _s3.exceptions.NoSuchKey:
pass
time.sleep(poll_interval_s)
raise InferenceTimeout(f"No result after {timeout_s}s — endpoint may still be cold-starting")
def extract_answer(result: dict) -> str:
truncated = False
if "choices" in result:
choice = result["choices"][0]
text = choice["message"]["content"]
truncated = choice.get("finish_reason") == "length"
elif "generated_text" in result:
text = result["generated_text"]
else:
return json.dumps(result)
if "<unused94>" in text and "<unused95>" not in text:
# Generation was cut off mid-reasoning before the real answer ever
# started -- nothing usable to show, so say so rather than dump the
# raw scratchpad.
return ("The model ran out of its response budget while reasoning "
"and never reached an answer. Try again, or ask a more "
"specific question.")
answer = _THINKING_RE.sub("", text).strip()
if truncated:
answer += ("\n\n*(Response was cut off — it hit the token limit "
"before finishing. Try asking a more specific question, "
"or ask it to continue.)*")
return answer
|