LongWriter: Unleashing 10,000+ Word Generation From Long Context LLMs

<p align="center"> 🤗 <a href="https://huggingface.co/datasets/THUDM/LongWriter-6k" target="_blank">HF Repo</a> • 📃 <a href="https://arxiv.org/abs/2408.07055" target="_blank">Paper</a> • 🚀 <a href="https://huggingface.co/spaces/THUDM/LongWriter" target="_blank">HF Space</a> </p>

English | 中文 | 日本語

https://github.com/user-attachments/assets/c7eedeca-98ed-43ec-8619-25137987bcde

Left: LongWriter-glm4-9b; Right: GLM-4-9B-chat

🔥 Updates

[2024/08/18] You can now deploy the LongWriter model using vllm. Refer to the code in vllm_inference.py and experience lightning-fast generation. It can generate over 10,000+ words in just one minute!

⚙️ LongWriter Deployment

Environmental Setup: We recommend using transformers>=4.43.0 to successfully deploy our models.

We open-source two models: LongWriter-glm4-9b and LongWriter-llama3.1-8b, trained based on GLM-4-9B and Meta-Llama-3.1-8B, respectively. These two models point to the "LongWriter-9B-DPO" and "LongWriter-8B" models in our paper. Try the model:

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
tokenizer = AutoTokenizer.from_pretrained("THUDM/LongWriter-glm4-9b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("THUDM/LongWriter-glm4-9b", torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
model = model.eval()
query = "Write a 10000-word China travel guide"
response, history = model.chat(tokenizer, query, history=[], max_new_tokens=32768, temperature=0.5)
print(response)

You may deploy your own LongWriter chatbot (like the one we show in the teasor video) by running

CUDA_VISIBLE_DEVICES=0 python trans_web_demo.py

Alternatively, you can deploy the model with vllm, which allows generating 10,000+ words within a minute! See the code example in vllm_inference.py.

🤖️ AgentWrite

agentwrite

We are also open-sourcing AgentWrite under agentwrite/, our automated ultra-long output data construction pipeline. Run plan.py and then write.py to obtain the final data. Please configure your API key in the files.

🖥️ Model Training

You can download and save the LongWriter-6k data through the Hugging Face datasets (🤗 HF Repo):

dataset = load_dataset('THUDM/LongWriter-6k')
for split, split_dataset in dataset.items():
    split_dataset.to_json("train/LongWriter-6k.jsonl")

You can mix it with your own general SFT data. We adopt the code and environment in LongAlign for model training (we use transformers==4.43.0 for training on Llama-3.1), with slight modification to adapt to new models. The training code is under train/. Please make sure to install FlashAttention 2 according to the code base of FlashAttention.

📊 Evaluation

We introduce two evaluation benchmarks: LongBench-Write and LongWrite-Ruler. LongBench-Write focuses more on measuring the long output quality as well as the output length, while LongWrite-Ruler is designed as a light-weight stress test of the model's maximum output length. We provide our evaluation data and code under evaluation/. Run

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python pred.py

to get model responses. Then run python eval_quality.py and python eval_length.py to evaluate the quality ($S_q$) and length ($S_l$) scores. Remember to configure your OpenAI API key in eval_quality.py since we adopt GPT-4o as the judge.

Here are the evaluation results on LongBench-Write: <img width="1000" alt="longbench-write" src="https://github.com/user-attachments/assets/8dbb6c02-09c4-4319-bd38-f1135457cd25"> Here are the evaluation results on LongWrite-Ruler:

👀 Cases

Here are LongWriter-glm4-9b's outputs to random test prompts.

User: Write a tragic love story about a lord's daughter falling in love with a servant, 5000 words.

<details> <summary>Assistant: (6176 words)</summary> <div style="max-height: 200px; overflow-y: auto; padding: 10px; border: 1px solid #e1e4e8; border-radius: 6px;">

The Aristocratic Estate

Nestled amidst the rolling hills of the English countryside, the grand estate of Lord Edward Ashwood was a beacon of elegance and sophistication. The sprawling manor, with its towering spires and expansive gardens, stood as a testament to centuries of wealth and power. The Ashwood family, known for their noble lineage and impeccable manners, had been pillars of the local community for generations. The grand halls were filled with the scent of polished wood and rich fabrics, echoing with the soft hum of servants attending to their masters' every need.

Lady Ashwood, the matriarch of the household, was a woman of grace and wisdom. Her beauty had faded with time, but her elegance remained untarnished. She was known for her charitable works and her ability to command respect and admiration from all who met her. Her daughters, Eliza and Arabella, were the epitome of aristocratic charm. Eliza, the elder, was a poised and intelligent young woman, often seen with a book in hand or engaged in scholarly discussions. Arabella, the younger, was vivacious and lively, her laughter ringing through the halls like a bell. Both girls were expected to marry well, their futures already mapped out by their parents.

The household was a bustling hive of activity. The grand dining hall was a testament to the family's wealth, adorned with intricate tapestries and fine china. The servants moved with practiced precision, their duties meticulously assigned. The butler, Mr. Blackwood, was a stern but fair man, ensuring that every aspect of the household ran smoothly. The head cook, Mrs. Brown, was a master of her craft, her culinary skills renowned throughout the county. The young page boys and maids scurried about, their faces a mix of innocence and the early signs of budding adulthood.

The Ashwood estate was a world apart from the simpler lives of the villagers who worked the surrounding lands. The gulf between the two worlds was vast and unbridgeable, a chasm that no one dared to cross. The servants, though integral to the estate's operation, were kept at a distance, their presence a mere background to the grandeur that defined the Ashwoods.

In this world of opulence and order, the seeds of a tragic love story were about to be sown.

The Hidden Heart of Eliza

Eliza Ashwood was the embodiment of aristocratic grace, her every movement a testament to years of refinement and education. Yet, beneath her composed exterior lay a heart yearning for something more, something the grand estate could not provide. Eliza was not merely a bookish daughter, though her love for literature was well-known. She possessed a depth of thought and a sensitivity that set her apart from her peers. Her days were filled with the routine of the estate—participating in social gatherings, learning the intricacies of diplomacy and etiquette, and engaging in scholarly pursuits. But her nights were often spent in quiet reflection, her mind wandering to thoughts that were far removed from the expectations of her station.

Eliza's interest in the estate's servants was a well-guarded secret. She often found herself drawn to the quieter corners of the manor, where the servants conducted their duties with a sense of quiet dignity. It was during one such moment, while taking a solitary walk through the estate's expansive gardens, that she first encountered Thomas. He was a young servant, tasked with tending to the roses that lined the estate's perimeter. His back was to her as he worked, his hands deftly pruning the delicate blooms. Eliza was captivated by the ease and precision with which he moved, a stark contrast to the hurried chaos of the household's daily operations.

As she watched him, a sense of tranquility washed over her. For the first time in what felt like an eternity, Eliza felt a sense of peace. Thomas's presence was unassuming, his focus entirely on his work. It was a stark contrast to the constant hustle and bustle of the manor. Eliza found herself lingering, her curiosity piqued by the man who seemed to exist in a world of his own, untouched by the grandeur that surrounded him.

Their first conversation was innocent and brief, a simple exchange over the health of Thomas's mother. Eliza's genuine concern and Thomas's humble gratitude left an impression on both of them. From that moment, their interactions became more frequent, though always under the watchful eyes of the household. Eliza would often find herself seeking out Thomas, her heart fluttering with a mix of excitement and fear. Their conversations were simple and heartfelt, often centered around the beauty of the garden or the weather. Yet, it was enough to kindle a flame in Eliza's heart, a flame that she knew she must keep hidden.

Eliza's feelings for Thomas were a source of both comfort and conflict. On one hand, he represented a connection to a simpler, more genuine world, a world where she felt truly seen and understood. On the other hand, she was acutely aware of the social and economic chasm that separated them. The idea of a relationship with a servant was not only unimaginable but also scandalous. Eliza's heart was a battlefield of hope and despair, torn between her desire for love and her duty to her family and society.

As their bond deepened, Eliza found herself questioning her place in the world. She had been raised to aspire to greatness, to marry a man of equal or greater status, to uphold the Ashwood legacy. Yet, here she was, falling in love with a servant, a man who represented everything her world sought to exclude. It was a love that defied logic and reason, a love that threatened to upend the carefully constructed life she had been destined to lead.

Eliza's secret affection for Thomas was a delicate balance, a fragile thread that she clung to with all her might. She knew that revealing her feelings would not only shatter her own dreams but also bring disgrace upon her family. Yet, the more time she spent with Thomas, the more she realized that her heart could no longer be contained. It was a love that was as beautiful as it was dangerous, a love that would test the very foundations of her world.

Thomas's Hidden Life

Thomas had been born into a family of humble means, his parents working the fields to make ends meet. From a young age, he had known the hardships of rural life, the backbreaking labor and the sting of poverty. Yet, it was not the struggle that defined him but the resilience and determination he found within himself. When the Ashwood estate advertised for new servants, Thomas saw it as an opportunity to escape the confines of his village and forge a better future for himself.

Life at the Ashwood estate was a stark contrast to his previous existence. The grandeur and opulence were overwhelming at first, but Thomas adapted quickly, his innate sense of duty driving him to excel in his duties. He was assigned to the garden, a role that suited his quiet nature and love for the natural world. Each day, he moved through the estate with a sense of purpose, his hands deftly tending to the plants, his mind a sanctuary of calm amidst the chaos.

Despite the differences in their worlds, Thomas felt a strange sense of belonging at the Ashwood estate. The grand halls and ornate decorations were a world apart from his simple upbringing, but there was a certain elegance to the place that he found comforting. The strict routines and the disciplined lives of the servants provided a structure that he had always craved. He found solace in the predictability of his tasks, in the rhythm of his workday that allowed him moments of quiet reflection.

It was during these moments of solitude that Thomas began to notice Eliza. Her presence was like a gentle breeze through the garden, her elegance contrasting with the rustic beauty of the flowers she admired. Thomas was initially drawn to her beauty, but as their interactions grew, he came to appreciate her intelligence and kindness. Eliza was unlike anyone he had ever met, her conversations filled with wisdom and compassion. She saw beyond his station, into the person he truly was.

Their bond grew slowly, built on shared moments of understanding and mutual respect. Thomas found himself looking forward to their brief conversations, his heart lightening with each encounter. Eliza's genuine interest in his life and her willingness to see past the societal barriers that separated them filled him with a sense of hope he had long thought lost. For the first time, Thomas felt valued and seen, not as a servant, but as a person with dreams and aspirations of his own.

Yet, Thomas was acutely aware of the dangers that their relationship posed. He knew the social and economic chasm that separated them, the impossibility of a future together. The idea of love was a fragile dream in a world that demanded conformity and respectability. Thomas's heart was a battlefield of hope and fear, torn between his desire for Eliza and the reality of their circumstances. He loved her with all his heart, but he was also pragmatic, understanding that their love was a risk he could not afford to take.

As their bond deepened, Thomas found himself questioning his place in the world. He had come to the Ashwood estate seeking a better life, but now he found himself caught in a web of emotions that threatened to unravel everything he had worked for. The love he felt for Eliza was a beautiful but dangerous distraction, a reminder of the dreams he dared not speak aloud. He knew that their love was a fragile thread, one that could easily be severed by the harsh realities of their world.

Thomas's heart was a sanctuary of love and fear, a place where his dreams and reality clashed. He loved Eliza with a passion that defied reason, but he was also realistic, understanding that their love was a fragile hope in a world that demanded conformity. As their bond grew stronger, Thomas found himself standing at a crossroads, his future uncertain and his heart in turmoil. He knew that his love for Eliza was a risk he was willing to take, but he was also aware of the dangers that lay ahead.

The Unspoken Bond

The bond between Eliza and Thomas grew stronger with each passing day, a silent yet powerful connection that neither could deny. Their conversations, though brief, were filled with a depth of understanding that transcended the barriers of their social standing. Eliza found herself looking forward to their encounters, each interaction a source of solace and joy in her otherwise structured and rigid world. Thomas, in turn, felt a sense of belonging and purpose that he had never known before, Eliza's presence a beacon of hope in his otherwise monotonous life.

One evening, as the sun dipped below the horizon, casting a golden glow over the garden, Eliza found herself once again drawn to the roses Thomas tended. This time, she approached him with a quiet determination. "Thomas," she began, her voice barely above a whisper, "I wanted to thank you for always being there, for listening to me when no one else would."

Thomas looked up, his eyes meeting hers with a mixture of surprise and gratitude. "It's my pleasure, Miss Eliza. You've always been kind to me."

Eliza took a deep breath, her heart pounding in her chest. "Thomas, there's something I need to tell you. I... I care for you deeply. More than just as a friend or a servant. I... I love you."

Thomas's eyes widened, his heart racing at the weight of her