astra-parser / app.py
Turhan123's picture
Update app.py
d17865e verified
Raw
History Blame Contribute Delete
1.7 kB
import os, json
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama
MODEL_PATH = os.path.join(os.environ.get("MODEL_DIR", "/home/user/models"),
os.environ.get("MODEL_FILE", "astra-meal-parser-q4_k_m.gguf"))
API_KEY = os.environ.get("API_KEY", "") # set as a Space secret
SYSTEM = (
"You are a meal parser. Extract every food item and its amount from the user's "
"meal description (Turkish or English). Return ONLY a strict JSON object of the form "
'{"items": [{"name": string, "amount": string}]}. '
"No macros, no calories, no conversational text, no markdown, only valid JSON."
)
llm = Llama(model_path=MODEL_PATH, n_ctx=2048, n_threads=2, chat_format="chatml")
app = FastAPI()
class MealIn(BaseModel):
meal: str
def extract_json(text):
s, e = text.find("{"), text.rfind("}")
if s == -1 or e == -1:
return None
try:
return json.loads(text[s:e+1])
except Exception:
return None
@app.get("/")
def health():
return {"status": "ok"}
@app.post("/parse")
def parse(body: MealIn, authorization: str = Header(default="")):
if API_KEY and authorization != f"Bearer {API_KEY}":
raise HTTPException(status_code=401, detail="Unauthorized")
out = llm.create_chat_completion(
messages=[{"role": "system", "content": SYSTEM},
{"role": "user", "content": body.meal}],
temperature=0, max_tokens=256, stop=["<|im_end|>"],
)
parsed = extract_json(out["choices"][0]["message"]["content"])
if parsed is None:
raise HTTPException(status_code=422, detail="Parse failed")
return parsed