Views
No views yet
1from transformers import AutoTokenizer, T5ForConditionalGeneration
2import json
3import re
4
5model = T5ForConditionalGeneration.from_pretrained('email_extractor_final_20251009_232152')
6tokenizer = AutoTokenizer.from_pretrained('email_extractor_final_20251009_232152')
7
8def extract_info(email_text):
9 input_text = f"extract company and role: {email_text}"
10 input_ids = tokenizer(input_text, return_tensors='pt', max_length=512, truncation=True).input_ids
11
12 outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
13 prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
14
15 # Fix JSON formatting
16 fixed = prediction.strip()
17 if fixed.startswith('"') and not fixed.startswith('{'):
18 fixed = '{' + fixed
19 if not fixed.endswith('}'):
20 fixed = fixed + '}'
21 fixed = re.sub(r'",(\s*)"', '", "', fixed)
22
23 return json.loads(fixed)
24
25# Example
26email = "Thank you for applying to OpenAI for the Software Engineer position."
27info = extract_info(email)
28print(f"Company: {info['company']}, Role: {info['role']}")