Потоковая передача результатов LLM в Django

Я планирую разработать Django фреймворк для LLM модели тонкой настройки на основе HuggingFace моделей, я хочу, чтобы текст был передан из модели в Django. Моя проблема в том, что модель передает вывод в терминал, но я не знаю, как сделать его в потоке в переменной, чтобы я мог передать его конечному пользователю

Я пробовал использовать конвейер с TextStream, но он выводит только текст в терминал


model_id = "model_a"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    # attn_implementation="flash_attention_2"
)
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True
)
streamer = TextStreamer(tokenizer)
pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    model_kwargs={"torch_dtype": torch.bfloat16},
    streamer=streamer
)

<<<В вашем случае, если вы замените 'TextStreamer' на TextIteratorStreamer, вы можете использовать стример следующим образом для потокового вывода.<<<0>> TextIteratorStreamer

>
    def get_response(question):
        query_response = YOUR_LLM_CHAIN(question)


    generation_thread = Thread(target=get_response, args=("Tell me a joke",))
    generation_thread.start()

    for new_token in streamer:
        yield new_token

При включенной потоковой передаче в виде текстового/событийного потока вы сможете передавать новые токены по мере их генерации.

Вернуться на верх