Tomiwajin/email-company-role-extractor

from transformers import AutoTokenizer, T5ForConditionalGeneration
import json
import re

model = T5ForConditionalGeneration.from_pretrained('email_extractor_final_20251009_232152')
tokenizer = AutoTokenizer.from_pretrained('email_extractor_final_20251009_232152')

def extract_info(email_text):
    input_text = f"extract company and role: {email_text}"
    input_ids = tokenizer(input_text, return_tensors='pt', max_length=512, truncation=True).input_ids
    
    outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
    prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Fix JSON formatting
    fixed = prediction.strip()
    if fixed.startswith('"') and not fixed.startswith('{'):
        fixed = '{' + fixed
    if not fixed.endswith('}'):
        fixed = fixed + '}'
    fixed = re.sub(r'",(\s*)"', '", "', fixed)
    
    return json.loads(fixed)

# Example
email = "Thank you for applying to OpenAI for the Software Engineer position."
info = extract_info(email)
print(f"Company: {info['company']}, Role: {info['role']}")