from typing import Dict, List, Any from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline, LongformerTokenizer import torch import torch.nn as nn import torch.nn.functional as F import spacy from spacy.matcher import PhraseMatcher from transformers import LongformerModel from skillNer.general_params import SKILL_DB from skillNer.skill_extractor_class import SkillExtractor Job_num_labels = None class EndpointHandler(): def __init__(self, path=""): # Label mapping as provided self.Job_label_map = { "JT": 0, # Job Title "JS": 1, # Job Summary "COT": 2, # Title of Company Overview Section "COC": 3, # Content of Company Overview Section "RT": 4, # Title of Responsibilites Section "RC": 5, # Content of Responsibilites Section "RQT": 6, # Title of Required Qualifications Section "RQC": 7, # Content of Required Qualifications Section "PQT": 8, # Title of Preferred Qualifications Section "PQC": 9, # Content of Preferred Qualifications Section "ET": 10, # Employment Type "SBC": 11, # Content of Salary and Benefits Section "SBT": 12 # Title of Salary and Benefits Section } global Job_num_labels self.Job_num_labels = len(self.Job_label_map) Job_num_labels = self.Job_num_labels self.Job_labels = [ {"value": "JT", "label": "Job Title"}, {"value": "JS", "label": "Job Summary"}, {"value": "COT", "label": "Title of Company Overview Section"}, {"value": "COC", "label": "Content of Company Overview Section"}, {"value": "RT", "label": "Title of Responsibilites Section"}, {"value": "RC", "label": "Content of Responsibilites Section"}, {"value": "RQT", "label": "Title of Required Qualifications Section"}, {"value": "RQC", "label": "Content of Required Qualifications Section"}, {"value": "PQT", "label": "Title of Preferred Qualifications Section"}, {"value": "PQC", "label": "Content of Preferred Qualifications Section"}, {"value": "ET", "label": "Employment Type"}, {"value": "SBC", "label": "Content of Salary and Benefits Section"}, {"value": "SBT", "label": "Title of Salary and Benefits Section"}, ] # Load tokenizer self.Job_tokenizer = LongformerTokenizer.from_pretrained("allenai/longformer-base-4096") self.Job_tokenizer.cls_token # Load model architecture self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.Job_model = LongformerSentenceClassifier(num_labels=self.Job_num_labels) self.Job_model.to(self.device) # Load trained weights self.Job_model.load_state_dict(torch.load(path + "/JobSegmentClassifier3rdEpoch_v2.pth", map_location=self.device)) # Set model to evaluation mode self.Job_model.eval() nlp = spacy.load("en_core_web_lg") self.skill_extractor = SkillExtractor(nlp, SKILL_DB, PhraseMatcher) def predict_job_sections(self, model, text, tokenizer, device): model.eval() # Tokenize text and get input tensors encoding = tokenizer( text, return_tensors="pt", truncation=True, padding="max_length", max_length=4096 ) input_ids = encoding["input_ids"].to(device) attention_mask = encoding["attention_mask"].to(device) # Identify `[CLS]` positions (assuming each sentence starts with `[CLS]`) cls_positions = (input_ids == tokenizer.cls_token_id).nonzero(as_tuple=True)[1] cls_positions = cls_positions.unsqueeze(0).to(device) # Shape: (1, num_sentences) # Create global attention mask (Longformer requires at least 1 global attention token) global_attention_mask = torch.zeros_like(input_ids) global_attention_mask[:, cls_positions] = 1 # Assign global attention to `[CLS]` tokens # Run the model with torch.no_grad(): logits = model( input_ids=input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask, cls_positions=cls_positions ) # Shape: (1, num_sentences, num_labels) logits = logits.squeeze(0) # Shape: (num_sentences, num_labels) probs = F.softmax(logits, dim=-1) # Convert logits to probabilities predictions = torch.argmax(probs, dim=-1) # Get predicted label indices return predictions.cpu().numpy() # Convert to NumPy array for easy use def extract_job_sections(self, text): lines = text.splitlines() lines = [line for line in text.splitlines() if line.strip()] text = lines concatenated_text = " ".join(f"{self.Job_tokenizer.cls_token} {sentence}" for sentence in text) predictions = self.predict_job_sections(self.Job_model, concatenated_text, self.Job_tokenizer, self.device) return predictions, text def extract_job_requirements(self, text): lines = text.splitlines() lines = [line for line in text.splitlines() if line.strip()] text = lines concatenated_text = " ".join(f"{self.Job_tokenizer.cls_token} {sentence}" for sentence in text) predictions = self.predict_job_sections(self.Job_model, concatenated_text, self.Job_tokenizer, self.device) requirements = [] for i, pred in enumerate(predictions): if self.Job_labels[pred]['value'] == "RQC" and i < len(lines): requirements.append(lines[i]) return requirements def label_job_post(self, text): lines = self.extract_job_requirements(text) response = { "requirements": [] } for item in lines: response["requirements"].append(item) response["skills"] = [] seen = set() if response["requirements"]: # Only process if we have requirements annotations = self.skill_extractor.annotate(" ".join(response["requirements"])) if 'results' in annotations and 'full_matches' in annotations['results']: for result in annotations['results']['full_matches']: # Standardizing the skill names skill_info = SKILL_DB.get(result["skill_id"], {}) skill_name = skill_info.get('skill_name', 'Unknown Skill') if skill_name not in seen: seen.add(skill_name) response["skills"].append({'name': skill_name, 'skill_id': result["skill_id"]}) if 'results' in annotations and 'ngram_scored' in annotations['results']: for result in annotations['results']['ngram_scored']: if result['score'] >= 1: # Standardizing the skill names skill_info = SKILL_DB.get(result["skill_id"], {}) skill_name = skill_info.get('skill_name', 'Unknown Skill') if skill_name not in seen: seen.add(skill_name) response["skills"].append({'name': skill_name, 'skill_id': result["skill_id"]}) return response def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: """ data args: inputs (:obj: `str` | `PIL.Image` | `np.array`) kwargs Return: A :obj:`list` | `dict`: will be serialized and returned """ text = data['inputs'] # predictions, text = self.extract_job_sections(text) # requirements = self.extract_job_requirements(text) label_job_post = self.label_job_post(text) return label_job_post class LongformerSentenceClassifier(nn.Module): def __init__(self, model_name="allenai/longformer-base-4096", num_labels=Job_num_labels): """ Custom Longformer model for sentence classification. Args: model_name (str): Hugging Face Longformer model. num_labels (int): Number of possible sentence labels. """ super(LongformerSentenceClassifier, self).__init__() self.longformer = LongformerModel.from_pretrained(model_name) self.classifier = nn.Linear(self.longformer.config.hidden_size, num_labels) def forward(self, input_ids, attention_mask, global_attention_mask, cls_positions): """ Forward pass for sentence classification. Args: input_ids (Tensor): Tokenized input IDs, shape (batch_size, max_length) attention_mask (Tensor): Attention mask, shape (batch_size, max_length) global_attention_mask (Tensor): Global attention mask, shape (batch_size, max_length) cls_positions (List[Tensor]): Indices of `[CLS]` tokens for each batch element. """ outputs = self.longformer( input_ids=input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask ) last_hidden_state = outputs.last_hidden_state cls_positions = cls_positions.view(input_ids.shape[0], -1) cls_embeddings = last_hidden_state.gather(1, cls_positions.unsqueeze(-1).expand(-1, -1, last_hidden_state.size(-1))) logits = self.classifier(cls_embeddings) return logits if __name__ == "__main__": # init handler my_handler = EndpointHandler(path=".") # prepare sample payload payload = {"inputs": """ We are seeking an experienced Accounting Specialist to join our team. The Accounting Specialist will be responsible for various financial tasks, including reconciling accounts, assist with accounts payable, preparing financial reports, and assisting the Controller. The ideal candidate will have a strong background in accounting principles and practices, as well as proficiency in Quickbooks accounting software, Excel and financial concepts. Responsibilities: - Perform general ledger reconciliation to ensure accuracy of financial data - Prepare and analyze financial reports, bank reconciliations and analysis. - Collaborate with internal teams to ensure compliance with accounting policies and procedures - Support financial audits by providing necessary documentation and information - Accounts Payable - Multiple Companies - GL Reconciliations - Prepare Weekly, Monthly and Quarterly Commission Reports - Daily Bank Deposits Skills: - Proficiency in accounting software QuickBooks Online - Strong knowledge of corporate finance principles and practices - Experience with general ledger reconciliation - Ability to understand concise financial reports - Strong analytical skills for financial analysis - Knowledge of financial auditing processes - Understanding of cash flow analysis - Solid grasp of financial concepts such as revenue recognition, depreciation, and accruals - Ability to manage multiple priorities and time effectively. Pay: $50,000 - $60,000 per year Benefits: 401(k) matching Dental insurance Health insurance Paid time off Vision insurance Experience: Microsoft Excel: 3 years (Required) QuickBooks Online: 3 years (Required) """} # holiday_payload = {"inputs": "Today is a though day"} # test the handler non_holiday_pred=my_handler(payload) # holiday_payload=my_handler(holiday_payload) # show results print(non_holiday_pred) # print("holiday_payload", holiday_payload)