Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import random
3
4# Load model and tokenizer
5tokenizer = AutoTokenizer.from_pretrained("EssentialAI/EAI-Distill-0.5b", trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained("EssentialAI/EAI-Distill-0.5b")
7
8def chunk_text(text, max_char_per_doc=30000):
9 if len(text) <= max_char_per_doc:
10 return text
11
12 chunk_size = max_char_per_doc // 3
13 start = text[:chunk_size]
14
15 middle_start = chunk_size
16 middle_end = len(text) - chunk_size
17
18 mid_point = random.randint(middle_start + chunk_size//2, middle_end - chunk_size//2)
19
20 middle = text[mid_point - chunk_size//2:mid_point + chunk_size//2]
21 end = text[-chunk_size:]
22 return f"[beginning]\n{start}\n[middle]\n{middle}\n[end]\n{end}"
23
24def classify_document(text):
25 chunked_text = chunk_text(text)
26
27 messages = [
28 {"role": "system", "content": "taxonomy"},
29 {"role": "user", "content": chunked_text},
30 ]
31
32 prompt = tokenizer.apply_chat_template(
33 messages,
34 tokenize=False,
35 add_generation_prompt=True
36 )
37
38 inputs = tokenizer(prompt, return_tensors="pt")
39 outputs = model.generate(**inputs, max_new_tokens=100)
40 return tokenizer.decode(outputs[0], skip_special_tokens=True)
41
42# Example usage
43document_text = "Your document content here..."
44classification = classify_document(document_text)
45print(classification){FDC primary},{FDC secondary or skip}
{Bloom cognitive process primary (1-6)},{Bloom cognitive process secondary (1-6) or skip}
{Bloom knowledge domain primary (1-4)},{Bloom knowledge domain secondary (1-4) or skip}
{Document type v1 primary (1-17)},{Document type v1 secondary (1-17) or skip}
{Extraction artifacts primary (0-4)},{Extraction artifacts secondary (0-4) or skip}
{Missing content primary (0-6)},{Missing content secondary (0-6) or skip}
{Document type v2 primary (1-25)},{Document type v2 secondary (1-25) or skip}
{Reasoning depth primary (1-6)},{Reasoning depth secondary (1-6) or skip}
{Technical correctness primary (1-6)},{Technical correctness secondary (1-6) or skip}
{Educational level primary (1-5)},{Educational level secondary (1-5) or skip}1@misc{ai2025essentialwebv1024ttokens,
2 title={Essential-Web v1.0: 24T tokens of organized web data},
3 author={Essential AI and : and Andrew Hojel and Michael Pust and Tim Romanski and Yash Vanjani and Ritvik Kapila and Mohit Parmar and Adarsh Chaluvaraju and Alok Tripathy and Anil Thomas and Ashish Tanwer and Darsh J Shah and Ishaan Shah and Karl Stratos and Khoi Nguyen and Kurt Smith and Michael Callahan and Peter Rushton and Philip Monk and Platon Mazarakis and Saad Jamal and Saurabh Srivastava and Somanshu Singla and Ashish Vaswani},
4 year={2025},
5 eprint={2506.14111},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2506.14111},
9}