File size: 6,596 Bytes
42bc504 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 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) |