pydantic

app/main.py

@@ -1,51 +1,21 @@
-
-from fastapi import FastAPI, Request, Form
+from fastapi import FastAPI
 from fastapi.responses import JSONResponse
+from pydantic import BaseModel  # ✅ FIX: Use BaseModel not model
 from app.model_loader import load_model
 import torch
 
+# Initialize app and model
 app = FastAPI()
 model, tokenizer = load_model()
 
-# @app.post("/predict")
-# async def predict(request: Request):
-#     data = await request.json()
-#     input_text = data.get("input", "")
-#     inputs = tokenizer(input_text, return_tensors="pt")
-#     with torch.no_grad():
-#         output = model.generate(
-#             **inputs,
-#             max_new_tokens=60,
-#             do_sample=False,
-#             temperature=0.3
-#         )
-#     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):
+# ✅ Define request body schema using Pydantic
+class InputText(BaseModel):
     input: str
 
+# 🚀 Inference endpoint
 @app.post("/predict")
 async def predict(input_text: InputText):
-    # Use a structured, instruction-style prompt
+    # Create prompt
     prompt = (
         "You are a neuroscience research assistant.\n"
         "Determine if the following abstract was modified by an AI model.\n"
@@ -53,8 +23,10 @@ async def predict(input_text: InputText):
         f"Abstract:\n{input_text.input}\n\nAnswer:"
     )
 
+    # Tokenize
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
+    # Generate
     with torch.no_grad():
         output = model.generate(
             **inputs,
@@ -64,10 +36,10 @@ async def predict(input_text: InputText):
             pad_token_id=tokenizer.eos_token_id,
         )
 
+    # Decode and extract answer
     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:
@@ -75,4 +47,4 @@ async def predict(input_text: InputText):
     else:
         answer = "unknown"
 
-    return {"output": response_text, "answer": answer}
+    return {"output": response_text, "answer": answer}