File size: 907 Bytes
fbac5b6
 
 
 
 
 
 
 
2fd4429
fbac5b6
2fd4429
 
 
 
 
 
 
 
fbac5b6
2fd4429
fbac5b6
2fd4429
 
fbac5b6
2fd4429
 
 
 
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
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