| import os |
| import torch |
| import spaces |
| import gradio as gr |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| MODEL_ID = os.getenv("MODEL_ID", "GnLOLot/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking") |
|
|
| model = None |
| tokenizer = None |
|
|
| @spaces.GPU |
| def chat_fn(message, history): |
| global model, tokenizer |
| if model is None: |
| print("Loading model...", flush=True) |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_ID, |
| torch_dtype=torch.bfloat16, |
| device_map="auto", |
| trust_remote_code=True |
| ) |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| print("Model loaded", flush=True) |
| messages = [] |
| for h in history: |
| messages.append({"role": "user", "content": h[0]}) |
| messages.append({"role": "assistant", "content": h[1]}) |
| messages.append({"role": "user", "content": message}) |
|
|
| prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) |
|
|
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=512, |
| temperature=0.7, |
| top_p=0.9, |
| do_sample=True, |
| pad_token_id=tokenizer.pad_token_id |
| ) |
|
|
| return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip() |
|
|
| with gr.Blocks(title="MiniCPM5-1B Chat") as demo: |
| gr.Markdown(f"# MiniCPM5-1B Chat\n**Model:** `{MODEL_ID}`\n\nPowered by ZeroGPU (free GPU)") |
| gr.ChatInterface( |
| fn=chat_fn, |
| title=None, |
| description="First request loads the model (~30s), subsequent calls are faster." |
| ) |
|
|
| demo.launch() |
|
|
