staarlight/Phi-3-mini-instruct-model

1import os2import torch3from fastapi import FastAPI4from pydantic import BaseModel5from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline6import logging7from starlette.requests import Request8from starlette.middleware.base import BaseHTTPMiddleware9 10# === FILTER AGAR ?logs=container TIDAK MUNCUL DI LOG ===11 12class IgnoreLogsMiddleware(BaseHTTPMiddleware):13    async def dispatch(self, request: Request, call_next):14        # Jika request-nya adalah polling container logs, matikan logging15        if request.url.query == "logs=container":16            logging.getLogger("uvicorn.access").disabled = True17        else:18            logging.getLogger("uvicorn.access").disabled = False19 20        return await call_next(request)21 22# Juga filter level logger agar lebih aman23class NoContainerLogsFilter(logging.Filter):24    def filter(self, record):25        return "logs=container" not in record.getMessage()26 27logging.getLogger("uvicorn.access").addFilter(NoContainerLogsFilter())28 29# === Konfigurasi dasar ===30HF_TOKEN = os.environ.get("HF_TOKEN")31MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"32 33_pipe = None34 35def get_model():36    global _pipe37    if _pipe is not None:38        return _pipe39 40    print("🚀 Memuat tokenizer dan model...")41 42    # Load tokenizer43    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)44 45    # Cek GPU46    if torch.cuda.is_available():47        device_map = "auto"48        dtype = torch.float1649    else:50        device_map = None  # CPU-only51        dtype = torch.float3252 53    # Load model54    model = AutoModelForCausalLM.from_pretrained(55        MODEL_ID,56        dtype=dtype,57        device_map=device_map,58        token=HF_TOKEN,59    )60 61    # Pastikan model sudah tidak di 'meta' device62    if any(param.device.type == "meta" for param in model.parameters()):63        raise RuntimeError("❌ Model masih di meta device! Tunggu Space siap atau periksa device_map.")64 65    # Buat pipeline66    _pipe = pipeline(67        "text-generation",68        model=model,69        tokenizer=tokenizer,70        device=0 if torch.cuda.is_available() else -1, 71    )72 73    print("✅ Model siap digunakan.")74    return _pipe75 76 77# === FastAPI App ===78app = FastAPI(title="Private Phi-3 mini 4K API", version="1.1")79pipe = get_model()80 81class PromptRequest(BaseModel):82    prompt: str83    max_new_tokens: int = 102484    temperature: float = 1e-585 86@app.post("/api/predict")87def api_predict(req: PromptRequest):88    try:89        if pipe.model.device.type == "meta":90            raise RuntimeError("Model belum siap, coba refresh Space.")91 92        result = pipe(93            req.prompt,94            max_new_tokens=req.max_new_tokens,95            do_sample=True,96            temperature=req.temperature,97            pad_token_id=pipe.tokenizer.eos_token_id,98        )99        # Hapus prompt dari hasil100        text = result[0]["generated_text"].replace(req.prompt, "").strip()101        return {"response": text}102    except Exception as e:103        return {"error": str(e)}104 105@app.get("/")106def home():107    return {"status": "ready", "message": f"Model {MODEL_ID} loaded successfully"}108 109if __name__ == "__main__":110    import uvicorn111    uvicorn.run(app, host="0.0.0.0", port=7860)