extractor.py

# from transformers import RobertaTokenizerFast, AutoModelForTokenClassification
# import re
# import torch
# from itertools import cycle

# tokenizer = RobertaTokenizerFast.from_pretrained("mrfirdauss/robert-base-finetuned-cv")
# model = AutoModelForTokenClassification.from_pretrained("mrfirdauss/robert-base-finetuned-cv")

# id2label = {0: 'O',
#      1: 'B-NAME',
#      3: 'B-NATION',
#      5: 'B-EMAIL',
#      7: 'B-URL',
#      9: 'B-CAMPUS',
#      11: 'B-MAJOR',
#      13: 'B-COMPANY',
#      15: 'B-DESIGNATION',
#      17: 'B-GPA',
#      19: 'B-PHONE NUMBER',
#      21: 'B-ACHIEVEMENT',
#      23: 'B-EXPERIENCES DESC',
#      25: 'B-SKILLS',
#      27: 'B-PROJECTS',
#      2: 'I-NAME',
#      4: 'I-NATION',
#      6: 'I-EMAIL',
#      8: 'I-URL',
#      10: 'I-CAMPUS',
#      12: 'I-MAJOR',
#      14: 'I-COMPANY',
#      16: 'I-DESIGNATION',
#      18: 'I-GPA',
#      20: 'I-PHONE NUMBER',
#      22: 'I-ACHIEVEMENT',
#      24: 'I-EXPERIENCES DESC',
#      26: 'I-SKILLS',
#      28: 'I-PROJECTS'}

# def merge_subwords(tokens, labels):
#     merged_tokens = []
#     merged_labels = []

#     current_token = ""
#     current_label = ""

#     for token, label in zip(tokens, labels):
#         if token.startswith("Ġ"):
#             if current_token:
#                 # Append the accumulated subwords as a new token and label
#                 merged_tokens.append(current_token)
#                 merged_labels.append(current_label)
#             # Start a new token and label
#             current_token = token[1:]  # Remove the 'Ġ'
#             current_label = label
#         else:
#             # Continue accumulating subwords into the current token
#             current_token += token

#     # Append the last token and label
#     if current_token:
#         merged_tokens.append(current_token)
#         merged_labels.append(current_label)

#     return merged_tokens, merged_labels

# def chunked_inference(text, tokenizer, model, max_length=512):
#     # Tokenize the text with truncation=False to get the full list of tokens
#     tok = re.findall(r'\w+|[^\w\s]', text, re.UNICODE)
#     tokens = tokenizer.tokenize(tok, is_split_into_words=True)
#     # Initialize containers for tokenized inputs
#     input_ids_chunks = []
#     # Create chunks of tokens that fit within the model's maximum input size
#     for i in range(0, len(tokens), max_length - 2):  # -2 accounts for special tokens [CLS] and [SEP]
#         chunk = tokens[i:i + max_length - 2]
#         # Encode the chunks. Add special tokens via the tokenizer
#         chunk_ids = tokenizer.convert_tokens_to_ids(chunk)
#         chunk_ids = tokenizer.build_inputs_with_special_tokens(chunk_ids)
#         input_ids_chunks.append(chunk_ids)

#     # Convert list of token ids into a tensor
#     input_ids_chunks = [torch.tensor(chunk_ids).unsqueeze(0) for chunk_ids in input_ids_chunks]

#     # Predictions container
#     predictions = []

#     # Process each chunk
#     for input_ids in input_ids_chunks:
#         attention_mask = torch.ones_like(input_ids)  # Create an attention mask for the inputs
#         output = model(input_ids, attention_mask=attention_mask)
#         logits = output[0] if isinstance(output, tuple) else output.logits
#         predictions_chunk = torch.argmax(logits, dim=-1).squeeze(0)
#         predictions.append(predictions_chunk[1:-1])

#     # Optionally, you can convert predictions to labels here
#     # Flatten the list of tensors into one long tensor for label mapping
#     predictions = torch.cat(predictions, dim=0)
#     predicted_labels = [id2label[pred.item()] for pred in predictions]
#     return merge_subwords(tokens,predicted_labels)

# def process_tokens(tokens, tag_prefix):
#     # Process tokens to extract entities based on the tag prefix
#     entities = []
#     current_entity = {}
#     for token, tag in tokens:
#         if tag.startswith('B-') and tag.endswith(tag_prefix):
#             # Start a new entity
#             if current_entity:
#                 # Append the current entity before starting a new one
#                 entities.append(current_entity)
#                 current_entity = {}
#             current_entity['text'] = token
#             current_entity['type'] = tag
#         elif tag.startswith('I-') and (('GPA') == tag_prefix or tag_prefix == ('URL')) and tag.endswith(tag_prefix) and current_entity:
#             current_entity['text'] += '' + token
#         elif tag.startswith('I-') and tag.endswith(tag_prefix) and current_entity:
#             # Continue the current entity
#             current_entity['text'] += ' ' + token
#     # Append the last entity if there is one
#     if current_entity:
#         entities.append(current_entity)
#     return entities

# def predict(text):
#     tokens, predictions = chunked_inference(text, tokenizer, model)
#     data = list(zip(tokens, predictions))
#     profile = {
#         "name": "",
#         "links": [],
#         "skills": [],
#         "experiences": [],
#         "educations": []
#     }
#     profile['name'] = ' '.join([t for t, p in data if p.endswith('NAME')])

#     for skills in process_tokens(data, 'SKILLS'):
#       profile['skills'].append(skills['text'])
#     #Links
#     for links in process_tokens(data, 'URL'):
#       profile['links'].append(links['text'])
#     # Process experiences and education
#     workzip = []
#     exp = process_tokens(data, 'EXPERIENCES DESC')
#     designation = process_tokens(data, 'DESIGNATION')
#     comp = process_tokens(data, 'COMPANY')
#     if len(exp) >= len (designation) and len(exp) >= len(comp):
#         workzip = zip(cycle(designation),cycle(comp),exp)
#     elif len(designation)>=len(comp):
#         workzip = zip((designation),cycle(comp),cycle(exp))
#     else:
#         workzip = zip(cycle(designation),(comp),cycle(exp))
#     for designation, company, experience_desc in workzip:
#         profile['experiences'].append({
#             "start": None,
#             "end": None,
#             "designation": designation['text'],
#             "company": company['text'],  # To be filled in similarly
#             "experience_description": experience_desc['text']  # To be filled in similarly
#         })
#     for major, gpa, campus in zip(process_tokens(data, 'MAJOR'), process_tokens(data, 'GPA'), process_tokens(data, 'CAMPUS')):
#         profile['educations'].append({
#             "start": None,
#             "end": None,
#             "major": major['text'],
#             "campus": campus['text'],  # To be filled in similarly
#             "GPA": gpa['text'] # To be filled in similarly
#         })

#     return profile