ArxAlfa/AIIASpace

1import torch2import torch.nn as nn3import torch.optim as optim4import numpy as np5from fastapi import FastAPI, UploadFile, File6from sklearn.metrics import mean_squared_error7import pandas as pd8from sklearn.model_selection import train_test_split9import csv10import io11 12 13# Define the DNN model14class DNN(nn.Module):15    def __init__(self, input_size, hidden_size, output_size, num_hidden_layers):16        super(DNN, self).__init__()17        self.fc1 = nn.Linear(input_size, hidden_size)18        self.relu1 = nn.ReLU()19        self.hidden_layers = nn.ModuleList()20        for _ in range(num_hidden_layers):21            self.hidden_layers.append(nn.Linear(hidden_size, hidden_size))22            self.hidden_layers.append(nn.ReLU())23        self.fc3 = nn.Linear(hidden_size, output_size)24 25    def forward(self, x):26        x = self.fc1(x)27        x = self.relu1(x)28        for layer in self.hidden_layers:29            x = layer(x)30        x = self.fc3(x)31        return x32 33 34# Load the model35model = DNN(input_size=6, hidden_size=64, output_size=1, num_hidden_layers=32)36 37device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")38model = model.to(device)39model.load_state_dict(torch.load("model_weights.pth", map_location=device))40 41# Create a new FastAPI app instance42app = FastAPI(docs_url="/", redoc_url="/new_redoc")43 44 45# Create a POST endpoint46@app.get(47    "/generate/{Soil_Quality}/{Seed_Variety}/{Fertilizer_Amount_kg_per_hectare}/{Sunny_Days}/{Rainfall_mm}/{Irrigation_Schedule}"48)49def generate(50    Soil_Quality: float,51    Seed_Variety: float,52    Fertilizer_Amount_kg_per_hectare: float,53    Sunny_Days: float,54    Rainfall_mm: float,55    Irrigation_Schedule: float,56):57    global model58 59    # Combine all inputs60    input_data = [61        Soil_Quality,62        Seed_Variety,63        Fertilizer_Amount_kg_per_hectare,64        Sunny_Days,65        Rainfall_mm,66        Irrigation_Schedule,67    ]68 69    input_data = torch.tensor([input_data], dtype=torch.float32)70    input_data = input_data.to(device)71    prediction = model(input_data)72    return {"prediction": prediction.item()}73 74    @app.post("/train")75    async def train(76        trainDatafile: UploadFile = File(...),77        testDatafile: UploadFile = File(...),78        epochs: int = 100,79    ):80        global model81 82        contents1 = await trainDatafile.read()83        train_data = pd.read_csv(io.StringIO(contents1.decode("utf-8")))84 85        contents2 = await testDatafile.read()86        test_data = pd.read_csv(io.StringIO(contents2.decode("utf-8")))87 88    # Convert data to numpy arrays89    X_train = train_data.drop("Yield_kg_per_hectare", axis=1).values90    y_train = train_data["Yield_kg_per_hectare"].values91    X_test = test_data.drop("Yield_kg_per_hectare", axis=1).values92    y_test = test_data["Yield_kg_per_hectare"].values93 94    # Convert data to torch tensors95    X_train = torch.tensor(X_train, dtype=torch.float32)96    X_train = X_train.to(device)97    y_train = torch.tensor(y_train, dtype=torch.float32)98    y_train = y_train.to(device)99 100    X_test = torch.tensor(X_test, dtype=torch.float32)101    X_test = X_test.to(device)102    y_test = torch.tensor(y_test, dtype=torch.float32)103 104    # Define loss function and optimizer105    criterion = nn.MSELoss()106    optimizer = optim.Adam(model.parameters(), lr=0.001)107 108    rmseList = []109 110    for epoch in range(epochs):111        optimizer.zero_grad()112 113        # Forward pass114        outputs = model(X_train)115        loss = criterion(outputs, y_train.unsqueeze(1))116 117        # Backward pass and optimization118        loss.backward()119        optimizer.step()120 121        predictions = model(X_test)122        rmse = np.sqrt(123            mean_squared_error(124                y_test.cpu().detach().numpy(), predictions.cpu().detach().numpy()125            )126        )127        print(128            f"Epoch: {epoch+1}, RMSE: {float(rmse)}, Loss: {float(np.sqrt(loss.cpu().detach().numpy()))}"129        )130        rmseList.append(float(rmse))131 132    torch.save(model.state_dict(), "model_weights.pth")133 134    return {"rmse": rmseList}135