HR_TOOL_V2 / modules /resume_parser.py
deepsu's picture
update-fix/Fixed major issues of code, api, and models. Model were updated with proper formatiing, FastAPI code is simpler and cleaner, finally the api is optimized by removing pdf parsings
f9a6bac
Raw History Blame Contribute Delete
1.27 kB
import torch
class ResumeParser:
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
def format_data(self, txt):
prompt = f"""<|im_start|>system
Extract and format resume: name | role | years | skills | core skills | experiences | education | certifications.
Extract ONLY from input. NO inference or additions.
<|im_end|>
<|im_start|>user
{txt}
<|im_end|>
<|im_start|>assistant
"""
return prompt
def parse(self, resume):
inputs = self.tokenizer(resume, return_tensors="pt", padding=False, truncation=True)
with torch.no_grad():
# Mixed precision for faster CPU inference
with torch.cpu.amp.autocast():
output = self.model.generate(
**inputs,
max_new_tokens=300,
do_sample=False,
use_cache=True # KV cache speeds up token generation
)
decoded = self.tokenizer.decode(output[0], skip_special_tokens=False)
decoded = decoded.split("<|im_start|>assistant\n")[1].split("\n<|im_end|>")[0]
return decoded