TIF_E41210014/api_svm.py

76 lines
2.1 KiB
Python

from flask import Flask, request, jsonify
import pickle
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
app = Flask(__name__)
# Load model dan TF-IDF
model_path = "model_svm.pkl"
tfidf_path = "tfidf_vectorizer.pkl"
def load_model():
with open(model_path, 'rb') as f:
model = pickle.load(f)
with open(tfidf_path, 'rb') as f:
vectorizer = pickle.load(f)
return model, vectorizer
model, vectorizer = load_model()
@app.route('/predict', methods=['POST'])
def predict():
data = request.json
full_text = data.get('full_text', '')
if not full_text:
return jsonify({"error": "full_text is required"}), 400
text_tfidf = vectorizer.transform([full_text])
prediction = model.predict(text_tfidf)[0]
return jsonify({"prediction": prediction})
@app.route('/train', methods=['POST'])
def train():
data = request.json
df = pd.DataFrame(data)
if 'full_text' not in df.columns or 'polarity' not in df.columns:
return jsonify({"error": "Data must contain 'full_text' and 'polarity'"}), 400
X = df['full_text']
y = df['polarity']
global vectorizer, model
vectorizer = TfidfVectorizer()
X_tfidf = vectorizer.fit_transform(X)
model = SVC()
model.fit(X_tfidf, y)
# Save model & vectorizer
with open(model_path, 'wb') as f:
pickle.dump(model, f)
with open(tfidf_path, 'wb') as f:
pickle.dump(vectorizer, f)
return jsonify({"message": "Model trained successfully"})
@app.route('/accuracy', methods=['POST'])
def accuracy():
data = request.json
df = pd.DataFrame(data)
if 'full_text' not in df.columns or 'polarity' not in df.columns:
return jsonify({"error": "Data must contain 'full_text' and 'polarity'"}), 400
X = df['full_text']
y_true = df['polarity']
X_tfidf = vectorizer.transform(X)
y_pred = model.predict(X_tfidf)
acc = accuracy_score(y_true, y_pred)
return jsonify({"accuracy": acc})
if __name__ == '__main__':
app.run(debug=True)