Update app.py

app.py

@@ -34,7 +34,7 @@ SYSTEM_PROMPT = (
 def root():
     return {"message": "MobileLLM-Pro API is running!"}
 
-@app.get("/generate")
+@app.get("/gen")
 def generate(prompt: str, max_tokens: int = 256):
     try:
         # Build messages with system instruction
@@ -70,5 +70,49 @@ def generate(prompt: str, max_tokens: int = 256):
 
         return {"input": prompt, "output": result.strip()}
 
+    except Exception as e:
+        return {"error": str(e)}
+
+@app.post("/generate")
+async def generate(request: Request):
+    try:
+        # Read JSON body from request
+        data = await request.json()
+        prompt = data.get("prompt", "")
+        max_tokens = data.get("max_tokens", 256)
+
+        # Build messages with system instruction
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": prompt}
+        ]
+
+        # Apply chat template
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            return_tensors="pt",
+            add_generation_prompt=True,
+            tokenize=True
+        ).to(device)
+
+        # Generate
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=inputs,
+                max_new_tokens=max_tokens,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+        # Decode only the new part
+        input_len = inputs.shape[1]
+        generated_tokens = outputs[0][input_len:]
+        result = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+        return {"input": prompt, "output": result.strip()}
+
     except Exception as e:
         return {"error": str(e)}