updated

README.md

@@ -1,12 +1,17 @@
 # Linear Regression Model
 
-This is a simple linear regression model that returns coefficients and intercept when provided with column names.
+This model provides two main functions:
+1. **Retrieve Coefficients and Intercept**: Return the model's coefficients and intercept based on input column names.
+2. **Predict Values**: Given column names and data values, return predictions.
 
 ## Usage
 
-### Input
-The model takes JSON input with the following structure:
+### 1. Retrieve Coefficients and Intercept
+Send a JSON payload with just column names:
+
 ```json
 {
-  "columns": ["feature1", "feature2", "feature3"]
+  "inputs": {
+    "columns": ["feature1", "feature2", "feature3"]
+  }
 }

model.py

@@ -1,22 +1,40 @@
 # model.py
 import joblib
+from typing import List, Dict, Union
+from transformers import Pipeline
 
 class LinearRegressionModel:
     def __init__(self):
-        # Load the model
+        # Load the trained linear regression model
         self.model = joblib.load("linear_regression_model.joblib")
 
-    def get_coefficients(self, columns):
-        # Return coefficients as a dictionary with column names as keys
+    def get_coefficients(self, columns: List[str]) -> Dict[str, float]:
+        # Return coefficients with column names as keys
         coefficients = dict(zip(columns, self.model.coef_.tolist()))
         intercept = self.model.intercept_
         return {"coefficients": coefficients, "intercept": intercept}
 
-# Instantiate the model for inference
+    def predict(self, columns: List[str], data: List[List[float]]) -> List[float]:
+        # Predict values for given input data
+        return self.model.predict(data).tolist()
+
+# Instantiate the model
 model = LinearRegressionModel()
 
-# Sample usage function to demonstrate
-def predict(inputs):
-    # Assume 'columns' is a key in the inputs dict with a list of column names
-    columns = inputs.get("columns", [])
-    return model.get_coefficients(columns)
+# Define a custom pipeline class for inference
+class CustomLinearRegressionPipeline(Pipeline):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def __call__(self, inputs: Dict[str, Union[List[str], List[List[float]]]]) -> Dict[str, Union[Dict[str, float], List[float]]]:
+        columns = inputs.get("columns", [])
+        data = inputs.get("data", [])
+        
+        if not data:  # If no data, return coefficients and intercept
+            return self.model.get_coefficients(columns)
+        else:  # If data is provided, return predictions
+            return {"predictions": self.model.predict(columns, data)}
+
+# Instantiate the custom pipeline
+pipeline = CustomLinearRegressionPipeline(model)

requirements.txt

@@ -0,0 +1,3 @@
+scikit-learn
+joblib
+transformers