Upload 2 files

flaskapp.py

@@ -0,0 +1,20 @@
+from flask import Flask, request, jsonify
+import torch
+
+from minilm_keyword_pred import predict_on_text
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+print('dev', device)
+
+app = Flask(__name__)
+
+@app.route('/get_keywords', methods=['POST'])
+def get_keywords():
+  text  = request.json.get('text', '')
+  if not text: return jsonify({"error": "The 'text' field is required."}), 400 
+  groups = predict_on_text(text)['predicted_groups_with_scores']
+  groups = [{'keyword': x[0], 'score': float(x[1])} for x in groups]
+  return jsonify(groups)
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=14005, debug=True, use_reloader=False)

minilm_keyword_pred.py

@@ -0,0 +1,151 @@
+import torch
+import torch.nn as nn
+import pickle
+from transformers import AutoTokenizer, AutoModel
+from tqdm import tqdm
+import numpy as np
+   
+OFFLINE_MODEL_PATH = "all-MiniLM-L6-v2" 
+
+# ==============================================================================
+#  STEP 1: DEFINE THE MODEL ARCHITECTURE
+#  This MUST be the exact same class definition you used for training.
+# ==============================================================================
+class ImprovedMultiTaskClassifier(nn.Module):
+    def __init__(self, model_name, num_keywords, num_groups, dropout_rate=0.1):
+        super(ImprovedMultiTaskClassifier, self).__init__()
+        self.transformer = AutoModel.from_pretrained(model_name)
+        hidden_size = self.transformer.config.hidden_size
+
+        self.keyword_classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.LayerNorm(hidden_size), nn.ReLU(), nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.LayerNorm(hidden_size // 2), nn.ReLU(), nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size // 2, num_keywords)
+        )
+        self.group_classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.LayerNorm(hidden_size), nn.ReLU(), nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.LayerNorm(hidden_size // 2), nn.ReLU(), nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size // 2, num_groups)
+        )
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
+        token_embeddings = outputs.last_hidden_state
+        attention_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        sum_embeddings = torch.sum(token_embeddings * attention_mask_expanded, 1)
+        sum_mask = torch.clamp(attention_mask_expanded.sum(1), min=1e-9)
+        pooled_output = sum_embeddings / sum_mask
+        keyword_logits = self.keyword_classifier(pooled_output)
+        group_logits = self.group_classifier(pooled_output)
+        return keyword_logits, group_logits
+
+
+# ==============================================================================
+#  STEP 2: LOAD ALL SAVED COMPONENTS
+# ==============================================================================
+print("Loading all components for inference...")
+
+# Set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Using device: {device}")
+
+# Load config
+with open('minilm_keyword_classifier_gemini/inference_config.pkl', 'rb') as f:
+    config = pickle.load(f)
+
+# *** IMPORTANT: Override the model_name to use the local path ***
+config['model_name'] = OFFLINE_MODEL_PATH
+# Load tokenizer from the same offline path
+tokenizer = AutoTokenizer.from_pretrained(OFFLINE_MODEL_PATH)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained('minilm_keyword_classifier_gemini/inference_tokenizer')
+
+# Load label encoders
+with open('minilm_keyword_classifier_gemini/inference_mlb_keywords.pkl', 'rb') as f:
+    mlb_keywords = pickle.load(f)
+with open('minilm_keyword_classifier_gemini/inference_mlb_groups.pkl', 'rb') as f:
+    mlb_groups = pickle.load(f)
+
+# Instantiate the model architecture
+num_keywords = len(mlb_keywords.classes_)
+num_groups = len(mlb_groups.classes_)
+model = ImprovedMultiTaskClassifier(config['model_name'], num_keywords, num_groups).to(device)
+
+# Load the trained weights
+model.load_state_dict(torch.load('minilm_keyword_classifier_gemini/inference_model.pth', map_location=device))
+
+# Set model to evaluation mode (very important!)
+model.eval()
+
+print("✅ All components loaded and model is ready for inference.")
+
+
+# ==============================================================================
+#  STEP 3: CREATE THE PREDICTION FUNCTION (MODIFIED TO INCLUDE SCORES)
+# ==============================================================================
+def predict_on_text(text: str):
+    """
+    Takes a string of text and returns the predicted keywords and groups
+    along with their confidence scores.
+    """
+    with torch.no_grad():
+        encoding = tokenizer(
+            text,
+            truncation=True,
+            padding='max_length',
+            max_length=512,
+            return_tensors='pt'
+        )
+        input_ids = encoding['input_ids'].to(device)
+        attention_mask = encoding['attention_mask'].to(device)
+
+        keyword_logits, group_logits = model(input_ids, attention_mask)
+
+        keyword_probs = torch.sigmoid(keyword_logits).cpu().numpy()[0]
+        group_probs = torch.sigmoid(group_logits).cpu().numpy()[0]
+
+        kw_threshold = config['optimal_keyword_threshold']
+        gr_threshold = config['optimal_group_threshold']
+
+        # --- MODIFICATION START ---
+        # Get keywords that are above the threshold
+        kw_indices = np.where(keyword_probs > kw_threshold)[0]
+        predicted_keywords_with_scores = [
+            (mlb_keywords.classes_[i], keyword_probs[i]) for i in kw_indices
+        ]
+
+        # Get groups that are above the threshold
+        gr_indices = np.where(group_probs > gr_threshold)[0]
+        predicted_groups_with_scores = [
+            (mlb_groups.classes_[i], group_probs[i]) for i in gr_indices
+        ]
+        
+        # Sort predictions by score in descending order
+        predicted_keywords_with_scores.sort(key=lambda x: x[1], reverse=True)
+        predicted_groups_with_scores.sort(key=lambda x: x[1], reverse=True)
+        # --- MODIFICATION END ---
+        
+        return {
+            'predicted_keywords_with_scores': predicted_keywords_with_scores,
+            'predicted_groups_with_scores': predicted_groups_with_scores,
+        }
+
+
+
+# list through all csv files in automarked\todo folder. Read the content column and loop through all the content there as text
+# for file in glob.glob('automarked\\todo\\*.csv'):
+#     with open(file, 'r', newline='', encoding='utf-8', errors='ignore') as f:
+#         reader = csv.DictReader(f)
+#         for row in reader:
+#             text = row['content']
+
+text = """I want you to understand, people think there are many problems in the world. There are no many problems in the world. There's only one problem in the world – human being. What other problem, I'm asking"""
+
+dpred = predict_on_text(text) 
+for d in dpred['predicted_groups_with_scores']:
+    print(d[0], d[1], d[1] > 0.5)