deploy-test / app.py
VortexHunter23's picture
Update app.py
2fd4429 verified
Raw
History Blame Contribute Delete
907 Bytes
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "Scalai/scal-lite-60b-code"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
_model_holder = {"model": None}
@spaces.GPU(duration=180) # generous — first load under MXFP4 can be slow
def _load_model():
if _model_holder["model"] is None:
_model_holder["model"] = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype="auto",
device_map="cuda",
trust_remote_code=True,
)
return _model_holder["model"]
# Force the load to happen under a REAL GPU context, not emulation
_load_model()
@spaces.GPU(duration=90)
def generate(prompt, system_prompt, max_new_tokens, temperature, top_p):
model = _model_holder["model"]
# ... rest of your generate function, using `model` instead of the old global