LongWriter: 長文コンテキストLLMsから10,000語以上の生成を解き放つ

🤗 <a href="https://huggingface.co/datasets/THUDM/LongWriter-6k" target="_blank">HF リポジトリ</a> • 📃 <a href="https://arxiv.org/abs/2408.07055" target="_blank">論文</a> • 🚀 <a href="https://huggingface.co/spaces/THUDM/LongWriter" target="_blank">HF スペース</a>

English | 中文 | 日本語

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

左: LongWriter-glm4-9b; 右: GLM-4-9B-chat

🔥 更新

[2024年8月18日] 現在、vllmを使用してLongWriterモデルをデプロイすることができます。コードはvllm_inference.pyに記載されており、超高速の生成を体験できます。1分以内に10,000字以上を生成することが可能です！

🔍 目次

• ⚙️ LongWriter デプロイ
• 🤖️ AgentWrite
• 🖥️ モデルトレーニング
• 📊 評価
• 👀 ケース
• 📝 引用

⚙️ LongWriter デプロイ

環境設定：モデルを正常にデプロイするために、transformers>=4.43.0 を使用することをお勧めします。

私たちは2つのモデルをオープンソース化しました：LongWriter-glm4-9b と LongWriter-llama3.1-8b、それぞれ GLM-4-9B と Meta-Llama-3.1-8B に基づいて訓練されています。これら2つのモデルは、論文中の「LongWriter-9B-DPO」と「LongWriter-8B」に対応しています。モデルを試してみてください：

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 = "中国旅行ガイドを10000語で書いてください"
response, history = model.chat(tokenizer, query, history=[], max_new_tokens=32768, temperature=0.5)
print(response)

以下のコマンドを実行して、独自のLongWriterチャットボットをデプロイすることができます（予告動画で紹介したものと同じように）：

CUDA_VISIBLE_DEVICES=0 python trans_web_demo.py

あるいは、vllmを使用してモデルをデプロイすることもできます。これにより、1分以内に10,000字以上を生成することが可能です！vllm_inference.pyのコード例を参照してください。

🤖️ AgentWrite

また、agentwrite/ 内に AgentWrite、超長出力データの自動生成パイプラインをオープンソース化しました。plan.py を実行し、その後 write.py を実行して最終データを取得してください。ファイル内でAPIキーを設定してください。

🖥️ モデルトレーニング

Hugging Faceデータセットを通じて LongWriter-6k データをダウンロードして保存することができます（🤗 HF リポジトリ）：

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

独自の一般的なSFTデータと混ぜて使用することができます。モデルトレーニングには LongAlign のコードと環境を採用しています（Llama-3.1 上でのトレーニングには transformers==4.43.0 を使用していますが）、新しいモデルに適応するためにわずかな変更を加えています。トレーニングコードは train/ ディレクトリ内にあります。 FlashAttention のコードベースに従って FlashAttention 2 をインストールするようにしてください。

📊 評価

私たちは、2つの評価基準を導入しました：LongBench-Write と LongWrite-Ruler。LongBench-Write は長い出力の質および出力の長さを測定することに重点を置き、LongWrite-Ruler はモデルの最大出力長さの軽量ストレステストとして設計されています。評価データとコードは evaluation/ ディレクトリ内に提供されています。以下のコマンドを実行して

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

モデルの応答を取得します。その後 python eval_quality.py と python eval_length.py を実行して、質（$S_q$）と長さ（$S_l$）のスコアを評価します。GPT-4oを審査として採用しているため、eval_quality.py 内に OpenAI APIキーを設定することを忘れないでください。

以下は LongBench-Write の評価結果です： 以下は LongWrite-Ruler の評価結果です：

👀 ケース

以下は、ランダムテストプロンプトに対する LongWriter-glm4-9b の出力です。

ユーザー：Write a tragic love story about a lord's daughter falling in love with a servant, 5000 words.

アシスタント：（6176語）

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 rem