基于Transformer实现GPT-J逐词/逐句文本生成的技术问询

可以实现类似ChatGPT的流式逐词生成，核心是利用transformers库的内置流式工具或手动循环生成单个token并实时输出。下面提供两种可行方案：

方案一：使用Hugging Face内置的TextStreamer #

这是最简单的实现方式，直接在model.generate中指定streamer参数，即可实时输出每个生成的token：

import transformers
import torch
from transformers import GPTJForCausalLM, AutoTokenizer, TextStreamer

# 初始化模型和tokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B", pad_token='<|endoftext|>', eos_token='<|endoftext|>', truncation_side='left')
model = GPTJForCausalLM.from_pretrained(
            "EleutherAI/gpt-j-6B",
            revision="float16",
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            use_cache=True,
            gradient_checkpointing=False  # 推理阶段关闭梯度检查点，提升生成速度
        )
model.to("cuda")

# 初始化流式输出器，跳过prompt和特殊token
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

prompt = "This is a test sentence, which should be completed"
inputs = tokenizer(prompt, return_tensors='pt').to("cuda")

# 生成时绑定streamer，实现实时输出
model.generate(**inputs,
               max_new_tokens=100,
               do_sample=True,
               top_k=15,
               top_p=0.9,
               temperature=1,
               no_repeat_ngram_size=4,
               pad_token_id=tokenizer.eos_token_id,
               streamer=streamer)

方案二：手动循环生成单个token（自定义性更强） #

如果需要自定义输出逻辑（比如逐句截断、添加格式），可以手动循环生成每个token，每次生成后追加到输入序列中：

import transformers
import torch
from transformers import GPTJForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B", pad_token='<|endoftext|>', eos_token='<|endoftext|>', truncation_side='left')
model = GPTJForCausalLM.from_pretrained(
            "EleutherAI/gpt-j-6B",
            revision="float16",
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            use_cache=True,
            gradient_checkpointing=False
        )
model.to("cuda")
model.eval()

prompt = "This is a test sentence, which should be completed"
inputs = tokenizer(prompt, return_tensors='pt').to("cuda")
generated_tokens = inputs.input_ids

# 循环生成最多100个新token
for _ in range(100):
    with torch.no_grad():
        outputs = model(generated_tokens)
        next_token_logits = outputs.logits[:, -1, :]
        
        # 应用和原代码一致的采样策略
        next_token_logits = transformers.top_k_top_p_filtering(next_token_logits, top_k=15, top_p=0.9)
        next_token = torch.multinomial(torch.softmax(next_token_logits / 1, dim=-1), num_samples=1)
        
        # 遇到结束符则停止生成
        if next_token.item() == tokenizer.eos_token_id:
            break
        
        # 将新token追加到序列中
        generated_tokens = torch.cat([generated_tokens, next_token], dim=-1)
        
        # 实时输出新增内容，避免重复打印prompt
        current_output = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
        print(current_output[len(prompt):], end='', flush=True)

print()  # 生成结束后换行

额外速度优化建议 #

• 关闭gradient_checkpointing：原代码开启的梯度检查点仅适用于训练，推理阶段会大幅拖慢速度，必须关闭。
• 保持use_cache=True：该参数会缓存之前的key/value特征，避免重复计算，是流式生成的核心优化点。
• 尝试4bit量化：使用bitsandbytes库对模型进行4bit量化，可进一步降低显存占用并提升生成速度。

内容的提问来源于stack exchange，提问作者BlackHawk