|
from sklearn.pipeline import Pipeline
|
|
from sklearn.naive_bayes import MultinomialNB
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from joblib import dump, load
|
|
import firebase_admin
|
|
from firebase_admin import credentials, firestore
|
|
import logging
|
|
import datetime
|
|
import re
|
|
import pandas as pd
|
|
import os
|
|
from flask import Flask, request, jsonify
|
|
|
|
|
|
logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
|
|
|
|
|
|
try:
|
|
cred_path = r"D:\app-sentinel-7qnr19-firebase-adminsdk-kjmbe-533749ec1a.json"
|
|
if not firebase_admin._apps:
|
|
cred = credentials.Certificate(cred_path)
|
|
firebase_admin.initialize_app(cred)
|
|
db = firestore.client()
|
|
logging.info("Firebase initialized successfully.")
|
|
except Exception as e:
|
|
logging.error(f"Error initializing Firebase: {e}")
|
|
db = None
|
|
|
|
|
|
try:
|
|
model_path = os.path.join(os.getcwd(), "model_pipeline.joblib")
|
|
model_pipeline = load(model_path)
|
|
logging.info("Model pipeline loaded successfully.")
|
|
except Exception as e:
|
|
logging.warning(f"Model pipeline not found. Training new one. Error: {e}")
|
|
try:
|
|
|
|
messages = ["example message 1", "example message 2"]
|
|
labels = ["label1", "label2"]
|
|
model_pipeline = Pipeline([
|
|
('vectorizer', TfidfVectorizer()),
|
|
('classifier', MultinomialNB())
|
|
])
|
|
model_pipeline.fit(messages, labels)
|
|
dump(model_pipeline, model_path)
|
|
logging.info("New model pipeline trained and saved.")
|
|
except Exception as e:
|
|
logging.error(f"Error training new model: {e}")
|
|
|
|
|
|
app = Flask(__name__)
|
|
|
|
@app.route('/classify', methods=['POST'])
|
|
def classify_message():
|
|
data = request.json
|
|
message = data.get("message", "")
|
|
|
|
|
|
if message.strip() == "":
|
|
return jsonify({"error": "Message cannot be empty"}), 400
|
|
|
|
classification = model_pipeline.predict([message])[0]
|
|
|
|
|
|
message_data = {
|
|
"text": message,
|
|
"classification": classification,
|
|
"timestamp": datetime.datetime.now(),
|
|
}
|
|
|
|
if db:
|
|
db.collection("all_messages").add(message_data)
|
|
logging.info(f"Message classified as {classification} and stored in Firestore.")
|
|
else:
|
|
logging.warning("Firestore is not initialized. Data not stored.")
|
|
|
|
return jsonify({"classification": classification})
|
|
|
|
if __name__ == "__main__":
|
|
app.run(debug=True)
|
|
|