File size: 1,980 Bytes
9d2ec8d 21b4ef9 7b4c18a 9d2ec8d 7b4c18a 9d2ec8d 7b4c18a 9d2ec8d 21b4ef9 7b4c18a 9d2ec8d 7b4c18a 21b4ef9 7b4c18a 21b4ef9 7b4c18a 21b4ef9 7b4c18a 21b4ef9 7b4c18a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# model.py
import joblib
import pandas as pd
from typing import List, Dict, Union
from transformers import Pipeline
class LinearRegressionModel:
def __init__(self):
# Load the trained linear regression model
self.model = joblib.load("linear_regression_model.joblib")
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}
def predict_from_csv(self, columns: List[str], csv_file_path: str) -> List[float]:
# Read the CSV file and select specified columns
data = pd.read_csv(csv_file_path)[columns]
return self.model.predict(data).tolist()
# Instantiate the model
model = LinearRegressionModel()
# 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], bytes]]) -> Dict[str, Union[Dict[str, float], List[float]]]:
columns = inputs.get("columns", [])
csv_file = inputs.get("csv_file", None) # Expect a CSV file in binary format
if not csv_file: # If no CSV file, return coefficients and intercept
return self.model.get_coefficients(columns)
else: # If CSV file is provided, save and process it for predictions
csv_file_path = "/tmp/input_data.csv" # Temporary path to save the CSV file
with open(csv_file_path, "wb") as f:
f.write(csv_file)
# Make predictions from CSV file
predictions = self.model.predict_from_csv(columns, csv_file_path)
return {"predictions": predictions}
# Instantiate the custom pipeline
pipeline = CustomLinearRegressionPipeline(model)
|