Потоковая передача результатов 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
При включенной потоковой передаче в виде текстового/событийного потока вы сможете передавать новые токены по мере их генерации.