AnoopUoI/simple-nn-api

1from fastapi import FastAPI2import torch3import torch.nn as nn4import torch.nn.functional as F5from pydantic import BaseModel6import uvicorn7 8app = FastAPI()9 10# Define a simple Neural Network11class SimpleNN(nn.Module):12    def __init__(self):13        super(SimpleNN, self).__init__()14        self.fc1 = nn.Linear(2, 16)15        self.fc2 = nn.Linear(16, 1)16 17    def forward(self, x):18        x = F.relu(self.fc1(x))19        x = torch.sigmoid(self.fc2(x))20        return x21 22# Initialize model23model = SimpleNN()24model.eval()  # Set model to evaluation mode25 26# Request Body Model27class InputData(BaseModel):28    x1: float29    x2: float30 31@app.post("/predict")32async def predict(data: InputData):33    # Convert input to tensor34    input_tensor = torch.tensor([[data.x1, data.x2]], dtype=torch.float32)35    with torch.no_grad():36        output = model(input_tensor)37        prediction = output.item()38    return {"prediction": prediction}39 40@app.get("/")41async def root():42    return {"message": "Simple Neural Network API is running!"}43 44if __name__ == "__main__":45    uvicorn.run(app, host="0.0.0.0", port=7860)46 47