dev2004v's picture
Update model.py
c369fec verified
import joblib
import os
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
# Define relative file paths (assuming 'models/' is in the same directory as this script)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "models", "random_forest_model.pkl")
VECTORIZER_PATH = os.path.join(BASE_DIR, "models", "vectorizer.pkl")
SVD_PATH = os.path.join(BASE_DIR, "models", "svd.pkl")
try:
# Load the trained model, vectorizer, and SVD transformer
model = joblib.load(MODEL_PATH)
vectorizer = joblib.load(VECTORIZER_PATH)
svd = joblib.load(SVD_PATH)
print(" Model, vectorizer, and SVD loaded successfully!")
except Exception as e:
print(f" Error loading model files: {e}")
exit(1)
def predict_text(text: str) -> dict:
"""Preprocess input text and predict using trained model."""
try:
X_tfidf = vectorizer.transform([text]) # Convert to TF-IDF
X_reduced = svd.transform(X_tfidf) # Apply dimensionality reduction
prediction = model.predict(X_reduced)[0] # Predict label (0 or 1)
probability = float(model.predict_proba(X_reduced)[0][1])
print(probability) # Printing probability
return {"generated": int(prediction), "probability": probability}
except Exception as e:
print(f" Prediction error: {e}")
return {"generated": -1, "probability": 0.0}