test_svc_model / app.py
yashriva's picture
Create app.py
41fdb00
raw
history blame
1.6 kB
from flask import Flask, jsonify, request
from transformers import pipeline, AutoTokenizer
import joblib
import json
# Load the model
model = joblib.load("iris_svm.joblib")
# Load the configuration file
with open("config.json", "r") as f:
config = json.load(f)
# Get the input features and target variable
features = config["features"]
target = config["targets"][0]
target_mapping = config["target_mapping"]
# Initialize the tokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
# Initialize the Flask app
app = Flask(__name__)
# Define the prediction route
@app.route("/predict", methods=["POST"])
def predict():
# Get the input data
input_data = request.json
# Construct the input text for the pipeline
input_text = f"SepalLengthCm: {input_data['SepalLengthCm']}, SepalWidthCm: {input_data['SepalWidthCm']}, PetalLengthCm: {input_data['PetalLengthCm']}, PetalWidthCm: {input_data['PetalWidthCm']}"
# Tokenize the input text
tokenized_input = tokenizer(input_text, return_tensors="pt")
# Make a prediction using the pipeline
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)
predicted_class_id = classifier(tokenized_input)[0]['label']
# Convert the predicted class ID to a class name
predicted_class_name = list(target_mapping.keys())[list(target_mapping.values()).index(predicted_class_id)]
# Return the predicted class name as a JSON response
return jsonify({"predicted_class": predicted_class_name})
# Run the Flask app
if __name__ == "__main__":
app.run(debug=True)