predict

app/main.py

@@ -1,5 +1,4 @@
-
-from fastapi import FastAPI, Request, Form
+from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 from app.model_loader import load_model
 import torch
@@ -11,8 +10,24 @@ model, tokenizer = load_model()
 async def predict(request: Request):
     data = await request.json()
     input_text = data.get("input", "")
-    inputs = tokenizer(input_text, return_tensors="pt")
+    
+    # Tokenize and move to model device
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    
+    # Generate next 15 tokens
     with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=50)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return JSONResponse(content={"output": response})
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=15,
+            do_sample=True,  # Optional: adds randomness
+            temperature=0.8,  # Optional: more natural output
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    # Decode only new part of generation
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Extract the continuation only (optional but useful)
+    continuation = generated_text[len(input_text):].strip()
+
+    return JSONResponse(content={"output": continuation})