handler.py

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":
                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": """
About the job
Job Title: Frontend Developer
Job Type: Full-time or Part-Time
Location: Remote

About Us:
Our mission at micro1 is to match the most talented people in the world with their dream jobs. If you are looking to be at the forefront of AI innovation and work with some of the fastest growing companies in Silicon Valley, we invite you to apply for a role. By joining the micro1 community, your resume will become visible to top industry leaders, unlocking access to the best career opportunities on the market.

Job Summary:
Join our dynamic team at micro1 as a Frontend Developer where you will be instrumental in creating engaging and dynamic user experiences for web applications. At micro1, we provide opportunities to work with leading companies in Silicon Valley, ensuring a stable and competitive income in a flexible remote setting. Embrace the opportunity to grow your career, access top industry opportunities, and enjoy a range of great benefits.

Key Responsibilities:
Develop high-quality frontend components for web applications.
Optimize user interfaces to enhance user experiences.
Collaborate with cross-functional teams to define and design new features.
Ensure cross-browser compatibility and implement responsive designs.
Maintain code quality and ensure responsiveness of applications.
Utilize industry best practices and design patterns.
Participate in code reviews to maintain high code quality standards.

Required Skills and Qualifications:
Proficiency in HTML, CSS, and JavaScript.
Strong experience with React and Angular frameworks.
Excellent written and verbal communication skills.
Ability to work effectively in a remote setting.
Demonstrated ability to develop and optimize user interfaces.
Experience with ensuring cross-browser compatibility of applications.
Strong problem-solving skills and attention to detail.

Preferred Qualifications:
Experience with responsive design and mobile-first development.
Familiarity with version control systems, such as Git.
Understanding of Agile methodologies.
It's okay if you don't. Having a Native to C2/C1 level in another language such as German, French, or Spanish is nice to have.
    """}
    # 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)