ludwig-ai / ludwig

Low-code framework for building custom LLMs, neural networks, and other AI models
http://ludwig.ai
Apache License 2.0
11.19k stars 1.19k forks source link

Add streaming support for zero shot inference #3878

Closed arnavgarg1 closed 10 months ago

arnavgarg1 commented 10 months ago

This PR introduces a new boolean flag called streaming to the LudwigModel.generate() API which allows users to see streaming output when performing zero-shot inference on single or multiple samples.

Demo

https://github.com/ludwig-ai/ludwig/assets/106701836/50fc01db-837b-4689-8e8a-8da491481da1

Config to reproduce the demo video

import yaml
import logging
from ludwig.api import LudwigModel

config = yaml.safe_load(
    """
model_type: llm
base_model: meta-llama/Llama-2-7b-chat-hf

quantization:
  bits: 4

input_features:
  - name: instruction
    type: text

output_features:
  - name: output
    type: text

generation:
  max_new_tokens: 64
  temperature: 0.1

trainer:
    type: none

backend:
  type: local
"""
)

model = LudwigModel(config, logging_level=logging.INFO)

# Single sample - Normal generation
output = model.generate("What is the meaning of life?", generation_config={"max_new_tokens": 32})

# Single sample - Streaming generation
output = model.generate("What is the meaning of life?", generation_config={"max_new_tokens": 32}, streaming=True)

# Multi sample - Normal generation
output = model.generate(["What is the meaning of life?", "What is the weather like in Germany around December?"], generation_config={"max_new_tokens": 64})

# Multi sample - Streaming generation
output = model.generate(["What is the meaning of life?", "What is the weather like in Germany around December?"], generation_config={"max_new_tokens": 64}, streaming=True)
github-actions[bot] commented 10 months ago

Unit Test Results

  6 files  ±0    6 suites  ±0   14m 10s :stopwatch: -5s 12 tests ±0    9 :heavy_check_mark: ±0    3 :zzz: ±0  0 :x: ±0  60 runs  ±0  42 :heavy_check_mark: ±0  18 :zzz: ±0  0 :x: ±0 

Results for commit caa80399. ± Comparison against base commit 22024d79.

:recycle: This comment has been updated with latest results.

arnavgarg1 commented 10 months ago

@alexsherstinsky I think I had a poor interface defined! Can you take a look now and see if it makes sense and your comments are addressed?