predict

app/main.py

@@ -22,20 +22,57 @@ model, tokenizer = load_model()
 #     response = tokenizer.decode(output[0], skip_special_tokens=True)
 #     return JSONResponse(content={"output": response})
 
+# @app.post("/predict")
+# async def predict(request: Request):
+#     data = await request.json()
+#     input_text = data.get("input", "")
+
+#     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+#     with torch.no_grad():
+#         outputs = model.generate(
+#             **inputs,
+#             max_new_tokens=120,
+#             do_sample=False,
+#             temperature=0.3
+#         )
+
+#     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+#     return JSONResponse(content={"output": response})
+
+class InputText(model):
+    input: str
+
 @app.post("/predict")
-async def predict(request: Request):
-    data = await request.json()
-    input_text = data.get("input", "")
+async def predict(input_text: InputText):
+    # Use a structured, instruction-style prompt
+    prompt = (
+        "You are a neuroscience research assistant.\n"
+        "Determine if the following abstract was modified by an AI model.\n"
+        "Respond with 'yes' if it was modified, or 'no' if it is original.\n\n"
+        f"Abstract:\n{input_text.input}\n\nAnswer:"
+    )
 
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
-        outputs = model.generate(
+        output = model.generate(
             **inputs,
-            max_new_tokens=120,
+            max_new_tokens=50,
             do_sample=False,
-            temperature=0.3
+            temperature=0.3,
+            pad_token_id=tokenizer.eos_token_id,
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return JSONResponse(content={"output": response})
+    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Extract only the model's answer
+    response_text = decoded_output[len(prompt):].strip().lower()
+    if "yes" in response_text:
+        answer = "yes"
+    elif "no" in response_text:
+        answer = "no"
+    else:
+        answer = "unknown"
+
+    return {"output": response_text, "answer": answer}