Upload NavyBayes.py

NavyBayes.py

@@ -1,5 +1,6 @@
-from sklearn.feature_extraction.text import TfidfVectorizer
+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
@@ -8,13 +9,13 @@ 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")
 
 # Firebase Initialization
 try:
-    # استخدم المسار الذي قدمته
     cred_path = r"D:\app-sentinel-7qnr19-firebase-adminsdk-kjmbe-533749ec1a.json"
     if not firebase_admin._apps:
         cred = credentials.Certificate(cred_path)
@@ -25,54 +26,55 @@ except Exception as e:
     logging.error(f"Error initializing Firebase: {e}")
     db = None
 
-# Load or Train Model
+# Load or Train Model with Pipeline
 try:
-    model_path = os.path.join(os.getcwd(), "model.joblib")
-    vectorizer_path = os.path.join(os.getcwd(), "vectorizer.joblib")
-    model = load(model_path)
-    vectorizer = load(vectorizer_path)
-    logging.info("Model and vectorizer loaded successfully.")
+    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 and vectorizer not found. Training new ones. Error: {e}")
-    # Train new model and vectorizer
-    messages = ["example message 1", "example message 2"]
-    labels = ["label1", "label2"]
-    vectorizer = TfidfVectorizer()
-    X = vectorizer.fit_transform(messages)
-    model = MultinomialNB()
-    model.fit(X, labels)
-    dump(model, model_path)
-    dump(vectorizer, vectorizer_path)
-    logging.info("New model and vectorizer trained and saved.")
-
-# Classify Message
-def classify_and_store_message(message):
-    """
-    Classify a message and store the result in Firestore.
-    """
-    global model, vectorizer
+    logging.warning(f"Model pipeline not found. Training new one. Error: {e}")
     try:
-        if not message.strip():
-            raise ValueError("Input message cannot be empty.")
+        # Train new model and vectorizer as part of pipeline
+        messages = ["example message 1", "example message 2"]
+        labels = ["label1", "label2"]
+        model_pipeline = Pipeline([
+            ('vectorizer', TfidfVectorizer()),  # تحويل النصوص إلى تمثيل رقمي
+            ('classifier', MultinomialNB())      # تصنيف باستخدام Naive Bayes
+        ])
+        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}")
 
-        # Transform the message using the vectorizer
-        message_vector = vectorizer.transform([message])
-        classification = model.predict(message_vector)[0]
+# Flask API
+app = Flask(__name__)
 
-        # Prepare data for Firestore
-        message_data = {
-            "text": message,
-            "classification": classification,
-            "timestamp": datetime.datetime.now(),
-        }
+@app.route('/classify', methods=['POST'])
+def classify_message():
+    data = request.json
+    message = data.get("message", "")
+    
+    # التصنيف باستخدام الـ pipeline
+    if message.strip() == "":
+        return jsonify({"error": "Message cannot be empty"}), 400
+    
+    classification = model_pipeline.predict([message])[0]
+    
+    # تخزين الرسالة في Firestore
+    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.")
-        else:
-            logging.warning("Firestore is not initialized. Data not stored.")
+    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})
 
-        return classification
-    except Exception as e:
-        logging.error(f"Error in classification: {e}")
-        return None
+if __name__ == "__main__":
+    app.run(debug=True)