Full MiniCPM5-1B ChatInterface with ZeroGPU

app.py

@@ -1,14 +1,56 @@
-import gradio as gr
-import spaces
+import os
 import torch
+import spaces
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-zero = torch.Tensor([0])
-print(f"zero on: {zero.device}", flush=True)
+MODEL_ID = os.getenv("MODEL_ID", "GnLOLot/MiniCPM5-1B-Claude-Opus-Fable5-V2-Thinking")
+
+model = None
+tokenizer = None
 
 @spaces.GPU
-def greet(n):
-    return f"Hello {zero + n} Tensor"
+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()
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
+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()