import os import json import joblib import torch import numpy as np from flask import Flask, request, render_template_string, jsonify from transformers import AutoTokenizer, AutoModelForSequenceClassification app = Flask(__name__) # ----------------------- # Load artifacts # ----------------------- SAVE_DIR = "./model" try: # Load model & tokenizer tokenizer = AutoTokenizer.from_pretrained(SAVE_DIR) model = AutoModelForSequenceClassification.from_pretrained(SAVE_DIR) model.eval() DEVICE = "cuda" if torch.cuda.is_available() else "cpu" model.to(DEVICE) # Load MultiLabelBinarizer and labels mlb = joblib.load(os.path.join(SAVE_DIR, "mlb.joblib")) with open(os.path.join(SAVE_DIR, "labels.json"), "r", encoding="utf-8") as f: labels = json.load(f) MODEL_LOADED = True print(f"Model loaded successfully on device: {DEVICE}") print(f"Available labels: {labels}") except Exception as e: MODEL_LOADED = False print(f"Error loading model: {e}") tokenizer = None model = None mlb = None labels = [] # Sigmoid for probabilities def sigmoid(x): return 1 / (1 + np.exp(-x)) # ----------------------- # Prediction function (single text only) # ----------------------- def predict_single(text, threshold=0.5): """Predict categories for a single text.""" if not MODEL_LOADED: return [], [] # Tokenize encodings = tokenizer( [text], # Wrap in list since model expects batch truncation=True, padding=True, max_length=256, return_tensors="pt" ).to(DEVICE) # Forward pass with torch.no_grad(): outputs = model(**encodings) logits = outputs.logits.cpu().numpy() # Convert to probabilities probs = sigmoid(logits) # Apply fixed threshold (0.5) pred_bin = (probs >= threshold).astype(int) # Decode to label names row_2d = np.array([pred_bin[0]]) categories = mlb.inverse_transform(row_2d)[0] return list(categories), probs[0] # HTML Template with embedded CSS + LinkedIn Footer HTML_TEMPLATE = """
Classify your customer review into relevant categories