Update app.py

app.py

@@ -1,66 +1,40 @@
-import torch
 import gradio as gr
+import numpy as np
+from sklearn.linear_model import LinearRegression
 import spaces
 
-# Define a simple model
-class SimpleAdditionModel(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.layer = torch.nn.Linear(2, 1)
-
-    def forward(self, x):
-        return self.layer(x)
-
-# Train and save the model
-def train_model():
-    model = SimpleAdditionModel()
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
-    criterion = torch.nn.MSELoss()
-
-    for _ in range(1000):
-        x = torch.rand(100, 2) * 10  # Random numbers between 0 and 10
-        y = x.sum(dim=1, keepdim=True)
+# Define the linear regression function
+def linear_regression(X, y, new_data):
+    X = np.array(X).reshape(-1, 1)  # Reshape for sklearn
+    y = np.array(y)
+
+    # Perform the computation within the GPU context
+    with spaces.GPU():
+        model = LinearRegression()
+        model.fit(X, y)
+        prediction = model.predict(np.array(new_data).reshape(-1, 1))
+
+    return prediction.tolist()
+
+# Create the Gradio interface
+with gr.Blocks(title="Linear Regression with ZeroGPU") as iface:
+    gr.Markdown("# Linear Regression Example")
+    gr.Markdown("This example performs linear regression using ZeroGPU for computation.")
+    
+    with gr.Row():
+        with gr.Column():
+            X = gr.Textbox(lines=5, label="Input Features (X)", placeholder="e.g., 1, 2, 3, 4")
+            y = gr.Textbox(lines=5, label="Target Values (y)", placeholder="e.g., 2, 4, 6, 8")
+            new_data = gr.Textbox(lines=1, label="New Data for Prediction", placeholder="e.g., 5")
+            submit_btn = gr.Button("Predict")
         
-        optimizer.zero_grad()
-        output = model(x)
-        loss = criterion(output, y)
-        loss.backward()
-        optimizer.step()
-
-    torch.save(model.state_dict(), 'simple_addition_model.pth')
-
-# Load the model
-@spaces.GPU
-def load_model():
-    model = SimpleAdditionModel()
-    model.load_state_dict(torch.load('simple_addition_model.pth'))
-    model.eval()
-    return model
-
-# Train and save the model if it doesn't exist
-try:
-    model = load_model()
-except FileNotFoundError:
-    print("Training model...")
-    train_model()
-    model = load_model()
-
-# Prediction function
-@spaces.GPU
-def add_numbers(num1, num2):
-    input_tensor = torch.tensor([num1, num2], dtype=torch.float32).unsqueeze(0)
-    with torch.no_grad():
-        result = model(input_tensor)
-    return round(result.item(), 2)
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=add_numbers,
-    inputs=[gr.Number(label="Number 1"), gr.Number(label="Number 2")],
-    outputs=gr.Number(label="Sum"),
-    title="Simple Addition Calculator",
-    description="Enter two numbers to see their sum (approximated by a neural network)."
-)
+        output = gr.Textbox(lines=5, label="Predicted Values")
+    
+    submit_btn.click(
+        fn=linear_regression,
+        inputs=[X, y, new_data],
+        outputs=output
+    )
 
 # Launch the interface
-iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch(server_name="0.0.0.0", server_port=7860)