From cfd09bfc47f0ad99097443e8f35f6027186e7cbf Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Tue, 7 Nov 2023 22:34:11 -0500
Subject: [PATCH] chore(runner): yield the outputs directly (#573)

update openai client examples to >1

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .pre-commit-config.yaml                       |   3 +
 bench.py                                      | 320 +++++++++++-------
 cz.py                                         |  18 +-
 examples/bentoml-demo/service.py              |   1 +
 examples/langchain-chains-demo/service.py     |  19 +-
 examples/langchain-tools-demo/service.py      |   1 +
 examples/openai_client.py                     |  69 ++--
 .../src/openllm_core/_configuration.py        |  25 +-
 openllm-core/src/openllm_core/_schemas.py     |  18 +-
 .../config/configuration_llama.py             |   1 -
 .../config/configuration_mistral.py           |   1 +
 openllm-python/pyproject.toml                 |   2 +-
 openllm-python/src/openllm/_llm.py            |   8 +-
 openllm-python/src/openllm/_runners.py        |   4 +-
 .../src/openllm/entrypoints/openai.py         |   2 +-
 .../src/openllm/playground/features.py        |   5 -
 tools/dependencies.py                         |   2 +-
 tools/update-config-stubs.py                  |   2 +-
 18 files changed, 311 insertions(+), 190 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d621ce4a..ffd0cb68 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -14,6 +14,9 @@ repos:
         verbose: true
         exclude: |
           (?x)^(
+              examples/.*|
+              bench.py |
+              cz.py |
               openllm-client/src/openllm_client/pb.*
           )$
   - repo: https://github.com/astral-sh/ruff-pre-commit
diff --git a/bench.py b/bench.py
index b9f71768..a2b38915 100755
--- a/bench.py
+++ b/bench.py
@@ -13,130 +13,222 @@ async def send_request(url, it, prompt, session, model, **attrs):
   config = openllm.AutoConfig.for_model(model).model_construct_env(**attrs).model_dump()
   data = {'prompt': prompt, 'llm_config': config, 'adapter_name': None}
   async with session.post(url, headers=headers, data=json.dumps(data)) as response:
-    result = await response.text()
-  print('-' * 10 + '\n\nreq:', it, ', prompt:', prompt, '\n\nGeneration:', result)
+    try:
+      result = await response.json()
+      print('-' * 10 + '\n\nreq:', it, ', prompt:', prompt, '\n\nGeneration:', result['outputs'][0]['text'])
+    except Exception as err:
+      print('Exception while sending request %d (%s):' % (it, prompt), await response.text(), err)
+
 
 async def main(args: argparse.Namespace) -> int:
   endpoint = 'generate' if args.generate else 'generate_stream'
   url = f'{os.getenv("OPENLLM_ENDPOINT", "http://localhost:3000")}/v1/{endpoint}'
   # len=572
   prompts = [
-      "Translate the following English text to French: 'Hello, how are you?'", "Summarize the plot of the book 'To Kill a Mockingbird.'",
-      'Generate a list of 10 random numbers between 1 and 100.', 'What is the capital of France?', 'Write a poem about nature.', 'Convert 25 degrees Celsius to Fahrenheit.',
-      'Describe the process of photosynthesis.', 'Tell me a joke.', 'List five famous scientists and their contributions to science.',
-      'Write a short story about a detective solving a mystery.', 'Explain the theory of relativity.', 'Provide a brief history of the Roman Empire.',
-      'Create a shopping list for a BBQ party.', "Write a movie review for the film 'Inception.'", 'Explain the concept of artificial intelligence.', 'Write a letter to your future self.',
-      'Describe the life cycle of a butterfly.', 'List the top 10 tourist destinations in Europe.', 'Explain the principles of supply and demand.',
-      'Create a menu for a vegetarian restaurant.', 'Write a haiku about the ocean.', 'Explain the importance of renewable energy sources.',
-      'List the ingredients for making chocolate chip cookies.', 'Write a persuasive essay on the benefits of exercise.', 'Describe the cultural significance of the Taj Mahal.',
-      'Explain the process of DNA replication.', 'Write a speech about the importance of education.', 'List the steps to start a small business.', 'Explain the concept of biodiversity.',
-      'Create a playlist for a road trip.', 'Write a short biography of Albert Einstein.', 'Describe the impact of social media on society.', 'Explain the principles of good nutrition.',
-      'List the 10 tallest mountains in the world.', 'Write a product review for a smartphone.', 'Create a workout routine for building muscle.', 'Explain the concept of climate change.',
-      'Describe the life and achievements of Marie Curie.', 'List the ingredients for making a classic margarita.', 'Write a blog post about time management.',
-      'Explain the process of cellular respiration.', 'Create a budget for a family vacation.', "Write a book summary for 'The Great Gatsby.'", 'Describe the history of the Internet.',
-      'Explain the principles of effective communication.', 'List the top 10 historical landmarks in the world.', 'Write a love letter to someone special.',
-      'Explain the concept of human rights.', 'Create a recipe for homemade pizza.', 'Write a movie script for a short film.', 'Describe the structure of the atom.',
-      'List the 10 most influential artists of the 20th century.', 'Explain the process of mitosis.', 'Create a travel itinerary for a trip to Japan.',
-      'Write a poem about the beauty of nature.', 'Explain the importance of environmental conservation.', 'List the essential items for a hiking trip.',
-      'Write a short story set in a post-apocalyptic world.', 'Describe the history of the Olympic Games.', 'Explain the principles of democracy.',
-      'Create a business plan for a tech startup.', 'Write a letter of recommendation for a colleague.', 'List the ingredients for a classic Caesar salad.',
-      'Explain the concept of artificial neural networks.', 'Describe the life and work of Leonardo da Vinci.', 'List the 10 most popular tourist attractions in the United States.',
-      'Write a persuasive speech on the dangers of smoking.', 'Explain the process of natural selection.', 'Create a menu for a fine dining restaurant.',
-      'Write a poem about the beauty of the night sky.', 'Explain the importance of renewable energy.', 'List the necessary equipment for a camping trip.',
-      'Write a short biography of William Shakespeare.', 'Describe the impact of social media on business marketing.', 'Explain the principles of project management.',
-      'Create a playlist for a relaxing evening at home.', 'Write a blog post about the history of space exploration.', 'Explain the process of protein synthesis.',
-      'List the 10 most famous landmarks in Europe.', 'Write a book review for a classic novel.', 'Describe the history of ancient Egypt.', 'Explain the concept of cultural diversity.',
-      'Create a recipe for a gourmet sandwich.', 'Write a screenplay for a science fiction movie.', "Describe the structure of the Earth's atmosphere.",
-      'List the 10 greatest inventions of all time.', 'Explain the process of meiosis.', 'Create a travel guide for a visit to Paris.', 'Write a poem about the changing seasons.',
-      'Explain the importance of clean energy sources.', 'List the essential camping gear for a wilderness adventure.', 'Write a short story about a time-traveling adventure.',
-      'Describe the history of the Renaissance.', 'Explain the principles of economics.', 'Create a business proposal for a new restaurant.',
-      'Write a letter to your future self 10 years from now.', 'List the ingredients for a classic lasagna.', 'Explain the concept of machine learning.',
-      'Describe the life and contributions of Martin Luther King Jr.', 'List the 10 most famous museums in the world.',
-      'Write a persuasive essay on the importance of environmental conservation.', 'Explain the process of geological erosion.', 'Create a menu for a vegan cafe.',
-      'Write a poem about the power of imagination.', 'Explain the significance of the Industrial Revolution.', 'List the items needed for a beach vacation.',
-      'Write a short biography of Charles Darwin.', 'Describe the impact of globalization on cultures.', 'Explain the principles of time management.',
-      'Create a playlist for a high-energy workout.', 'Write a blog post about the future of artificial intelligence.', 'Explain the process of DNA transcription.',
-      'List the 10 most iconic landmarks in Asia.', 'Write a book summary for a popular self-help book.', 'Describe the history of the ancient Greeks.',
-      'Explain the concept of social justice.', 'Create a recipe for a gourmet salad.', 'Write a screenplay for a romantic comedy movie.', "Describe the layers of the Earth's atmosphere.",
-      'List the 10 most influential inventors in history.', 'Explain the process of plate tectonics.', 'Create a travel itinerary for a road trip across the USA.',
-      'Write a poem about the wonders of the natural world.', 'Explain the importance of sustainable agriculture.', 'List the essential hiking gear for a mountain expedition.',
-      'Write a short story about a futuristic dystopia.', 'Describe the history of the Middle Ages.',
-      'Write a letter to your future self, offering reflections on personal growth, achievements, and aspirations, as well as words of encouragement and guidance for your future journey.',
-      'List the ingredients for a classic chicken pot pie recipe, a beloved comfort food that combines tender chicken, vegetables, and a flaky pastry crust in a savory filling.',
-      'Explain the concept of artificial neural networks and their pivotal role in machine learning and artificial intelligence applications, from image recognition to natural language processing.',
-      'Describe the life and contributions of Albert Einstein, shedding light on his groundbreaking theories of relativity, his influence on modern physics, and his enduring legacy.',
-      "List the 10 most iconic landmarks in Australia, celebrating the country's breathtaking natural landscapes, architectural wonders, and unique cultural sites.",
-      'Write a persuasive speech on the importance of mental health awareness, advocating for destigmatization, access to mental health services, and compassionate support for those in need.',
-      "Explain the process of plate tectonics and its profound impact on the Earth's geology, including the formation of continents, mountain ranges, and the movement of tectonic plates.",
-      'Create a menu for a Mexican street food restaurant, featuring authentic and flavorful dishes that capture the essence of Mexican cuisine, from tacos to tamales.',
-      'Write a poem about the beauty of a moonlit night, exploring themes of serenity, reflection, and the enchanting ambiance created by the gentle glow of the moon.',
-      'Explain the significance of the Civil Rights Movement in the United States, highlighting the courageous individuals and leaders who paved the way for racial equality and justice.',
-      'List the items needed for a camping trip in the wilderness, ensuring that outdoor enthusiasts are well-equipped for their adventure in the great outdoors.',
-      'Write a short biography of Jane Austen, delving into her literary contributions and the enduring impact of her novels on literature, feminism, and societal norms.',
-      'Describe the impact of social media on political activism, discussing its role in mobilizing movements, shaping public discourse, and bringing about political change.',
-      'Explain the principles of effective leadership and management in the context of modern organizations, emphasizing the qualities and skills that define successful leaders.',
-      'Create a playlist for a high-energy workout session, curating motivating songs that provide the perfect soundtrack for a productive fitness routine.',
-      'Write a blog post about the future of artificial intelligence, exploring its potential applications in healthcare, transportation, and various industries that stand to be transformed.',
-      'Explain the process of DNA replication, shedding light on the intricate mechanism by which genetic information is faithfully duplicated during cell division.',
-      'List the 10 most famous landmarks in Asia, offering insights into their historical and cultural significance as well as their unique architectural features.',
-      'Write a book review for a contemporary novel that captivated your imagination, sharing your thoughts on the plot, characters, and overall impact of the book.',
-      'Describe the history of ancient Egypt, offering a glimpse into the rich civilization that thrived along the banks of the Nile River, from monumental architecture to hieroglyphic writing.',
-      'Explain the concept of cultural diversity, emphasizing the importance of embracing and celebrating diverse cultures and perspectives in a globalized world.',
-      'Create a recipe for a gourmet pasta dish that tantalizes the taste buds with a harmonious blend of flavors and textures, elevating a classic dish to culinary excellence.',
-      'Write a screenplay for a thrilling action movie set in a dystopian future, weaving together elements of suspense, adventure, and compelling characters in a post-apocalyptic world.',
-      "Describe the structure of the Earth's inner core, exploring its composition and the geophysical processes that drive the Earth's magnetic field and seismic activity.",
-      'List the 10 most significant technological innovations of the 21st century, from transformative advances in communication to breakthroughs in medical science and beyond.',
-      "Explain the process of geological earthquakes, delving into the geological forces and factors that contribute to the occurrence of seismic events and their impact on the Earth's surface.",
-      'Create a travel itinerary for a trip to Tokyo, Japan, offering recommendations for cultural experiences, sightseeing, and culinary delights in this vibrant metropolis.',
-      'Write a poem about the beauty of a sunset over the ocean, capturing the awe-inspiring colors and the sense of tranquility that descends as the sun dips below the horizon.',
-      'Explain the importance of conserving endangered species, highlighting the critical role these species play in maintaining ecosystem balance and the urgent need for conservation efforts.',
-      "List the essential gear for a camping and hiking adventure in a pristine natural wilderness, ensuring that outdoor enthusiasts are well-prepared for their journey into nature's beauty.",
-      'Write a short story about a group of explorers embarking on a thrilling quest to uncover a hidden treasure, navigating treacherous landscapes and solving ancient riddles along the way.',
-      'Describe the history of the Byzantine Empire, exploring its cultural achievements, architectural marvels such as the Hagia Sophia, and its enduring influence on Eastern Europe and beyond.',
-      'Explain the principles of ethical leadership, delving into the moral responsibilities of leaders in various domains and the impact of ethical leadership on organizations and society.',
-      'Create a business proposal for a sustainable fashion brand committed to eco-friendly practices, ethical sourcing, and transparency in the fashion industry.',
-      'Write a letter to your future self, offering reflections on personal growth, achievements, and aspirations, as well as words of encouragement and guidance for your future journey.',
-      'List the ingredients for a classic chicken pot pie recipe, a beloved comfort food that combines tender chicken, vegetables, and a flaky pastry crust in a savory filling.',
-      'Explain the concept of artificial neural networks and their pivotal role in machine learning and artificial intelligence applications, from image recognition to natural language processing.',
-      'Describe the life and contributions of Albert Einstein, shedding light on his groundbreaking theories of relativity, his influence on modern physics, and his enduring legacy.',
-      "List the 10 most iconic landmarks in Australia, celebrating the country's breathtaking natural landscapes, architectural wonders, and unique cultural sites.",
-      'Write a persuasive speech on the importance of mental health awareness, advocating for destigmatization, access to mental health services, and compassionate support for those in need.',
-      "Explain the process of plate tectonics and its profound impact on the Earth's geology, including the formation of continents, mountain ranges, and the movement of tectonic plates.",
-      'Create a menu for a Mexican street food restaurant, featuring authentic and flavorful dishes that capture the essence of Mexican cuisine, from tacos to tamales.',
-      'Write a poem about the beauty of a moonlit night, exploring themes of serenity, reflection, and the enchanting ambiance created by the gentle glow of the moon.',
-      'Explain the significance of the Civil Rights Movement in the United States, highlighting the courageous individuals and leaders who paved the way for racial equality and justice.',
-      'List the items needed for a camping trip in the wilderness, ensuring that outdoor enthusiasts are well-equipped for their adventure in the great outdoors.',
-      'Write a short biography of Jane Austen, delving into her literary contributions and the enduring impact of her novels on literature, feminism, and societal norms.',
-      'Describe the impact of social media on political activism, discussing its role in mobilizing movements, shaping public discourse, and bringing about political change.',
-      'Explain the principles of effective leadership and management in the context of modern organizations, emphasizing the qualities and skills that define successful leaders.',
-      'Create a playlist for a high-energy workout session, curating motivating songs that provide the perfect soundtrack for a productive fitness routine.',
-      'Write a blog post about the future of artificial intelligence, exploring its potential applications in healthcare, transportation, and various industries that stand to be transformed.',
-      'Explain the process of DNA replication, shedding light on the intricate mechanism by which genetic information is faithfully duplicated during cell division.',
-      'List the 10 most famous landmarks in Asia, offering insights into their historical and cultural significance as well as their unique architectural features.',
-      'Write a book review for a contemporary novel that captivated your imagination, sharing your thoughts on the plot, characters, and overall impact of the book.',
-      'Describe the history of ancient Egypt, offering a glimpse into the rich civilization that thrived along the banks of the Nile River, from monumental architecture to hieroglyphic writing.',
-      'Explain the concept of cultural diversity, emphasizing the importance of embracing and celebrating diverse cultures and perspectives in a globalized world.',
-      'Create a recipe for a gourmet pasta dish that tantalizes the taste buds with a harmonious blend of flavors and textures, elevating a classic dish to culinary excellence.',
-      'Write a screenplay for a thrilling action movie set in a dystopian future, weaving together elements of suspense, adventure, and compelling characters in a post-apocalyptic world.',
-      "Describe the structure of the Earth's inner core, exploring its composition and the geophysical processes that drive the Earth's magnetic field and seismic activity.",
-      'List the 10 most significant technological innovations of the 21st century, from transformative advances in communication to breakthroughs in medical science and beyond.',
-      "Explain the process of geological earthquakes, delving into the geological forces and factors that contribute to the occurrence of seismic events and their impact on the Earth's surface.",
-      'Create a travel itinerary for a trip to Tokyo, Japan, offering recommendations for cultural experiences, sightseeing, and culinary delights in this vibrant metropolis.',
-      'Write a poem about the beauty of a sunset over the ocean, capturing the awe-inspiring colors and the sense of tranquility that descends as the sun dips below the horizon.',
-      'Explain the importance of conserving endangered species, highlighting the critical role these species play in maintaining ecosystem balance and the urgent need for conservation efforts.',
-      "List the essential gear for a camping and hiking adventure in a pristine natural wilderness, ensuring that outdoor enthusiasts are well-prepared for their journey into nature's beauty.",
-      'Write a short story about a group of explorers embarking on a thrilling quest to uncover a hidden treasure, navigating treacherous landscapes and solving ancient riddles along the way.',
-      'Describe the history of the Byzantine Empire, exploring its cultural achievements, architectural marvels such as the Hagia Sophia, and its enduring influence on Eastern Europe and beyond.',
-      'Explain the principles of ethical leadership, delving into the moral responsibilities of leaders in various domains and the impact of ethical leadership on organizations and society.',
-      'Create a business proposal for a sustainable fashion brand committed to eco-friendly practices, ethical sourcing, and transparency in the fashion industry.',
-      'Write a letter to your future self, offering reflections on personal growth, achievements, and aspirations, as well as words of encouragement and guidance for your future journey.'
+    "Translate the following English text to French: 'Hello, how are you?'",
+    "Summarize the plot of the book 'To Kill a Mockingbird.'",
+    'Generate a list of 10 random numbers between 1 and 100.',
+    'What is the capital of France?',
+    'Write a poem about nature.',
+    'Convert 25 degrees Celsius to Fahrenheit.',
+    'Describe the process of photosynthesis.',
+    'Tell me a joke.',
+    'List five famous scientists and their contributions to science.',
+    'Write a short story about a detective solving a mystery.',
+    'Explain the theory of relativity.',
+    'Provide a brief history of the Roman Empire.',
+    'Create a shopping list for a BBQ party.',
+    "Write a movie review for the film 'Inception.'",
+    'Explain the concept of artificial intelligence.',
+    'Write a letter to your future self.',
+    'Describe the life cycle of a butterfly.',
+    'List the top 10 tourist destinations in Europe.',
+    'Explain the principles of supply and demand.',
+    'Create a menu for a vegetarian restaurant.',
+    'Write a haiku about the ocean.',
+    'Explain the importance of renewable energy sources.',
+    'List the ingredients for making chocolate chip cookies.',
+    'Write a persuasive essay on the benefits of exercise.',
+    'Describe the cultural significance of the Taj Mahal.',
+    'Explain the process of DNA replication.',
+    'Write a speech about the importance of education.',
+    'List the steps to start a small business.',
+    'Explain the concept of biodiversity.',
+    'Create a playlist for a road trip.',
+    'Write a short biography of Albert Einstein.',
+    'Describe the impact of social media on society.',
+    'Explain the principles of good nutrition.',
+    'List the 10 tallest mountains in the world.',
+    'Write a product review for a smartphone.',
+    'Create a workout routine for building muscle.',
+    'Explain the concept of climate change.',
+    'Describe the life and achievements of Marie Curie.',
+    'List the ingredients for making a classic margarita.',
+    'Write a blog post about time management.',
+    'Explain the process of cellular respiration.',
+    'Create a budget for a family vacation.',
+    "Write a book summary for 'The Great Gatsby.'",
+    'Describe the history of the Internet.',
+    'Explain the principles of effective communication.',
+    'List the top 10 historical landmarks in the world.',
+    'Write a love letter to someone special.',
+    'Explain the concept of human rights.',
+    'Create a recipe for homemade pizza.',
+    'Write a movie script for a short film.',
+    'Describe the structure of the atom.',
+    'List the 10 most influential artists of the 20th century.',
+    'Explain the process of mitosis.',
+    'Create a travel itinerary for a trip to Japan.',
+    'Write a poem about the beauty of nature.',
+    'Explain the importance of environmental conservation.',
+    'List the essential items for a hiking trip.',
+    'Write a short story set in a post-apocalyptic world.',
+    'Describe the history of the Olympic Games.',
+    'Explain the principles of democracy.',
+    'Create a business plan for a tech startup.',
+    'Write a letter of recommendation for a colleague.',
+    'List the ingredients for a classic Caesar salad.',
+    'Explain the concept of artificial neural networks.',
+    'Describe the life and work of Leonardo da Vinci.',
+    'List the 10 most popular tourist attractions in the United States.',
+    'Write a persuasive speech on the dangers of smoking.',
+    'Explain the process of natural selection.',
+    'Create a menu for a fine dining restaurant.',
+    'Write a poem about the beauty of the night sky.',
+    'Explain the importance of renewable energy.',
+    'List the necessary equipment for a camping trip.',
+    'Write a short biography of William Shakespeare.',
+    'Describe the impact of social media on business marketing.',
+    'Explain the principles of project management.',
+    'Create a playlist for a relaxing evening at home.',
+    'Write a blog post about the history of space exploration.',
+    'Explain the process of protein synthesis.',
+    'List the 10 most famous landmarks in Europe.',
+    'Write a book review for a classic novel.',
+    'Describe the history of ancient Egypt.',
+    'Explain the concept of cultural diversity.',
+    'Create a recipe for a gourmet sandwich.',
+    'Write a screenplay for a science fiction movie.',
+    "Describe the structure of the Earth's atmosphere.",
+    'List the 10 greatest inventions of all time.',
+    'Explain the process of meiosis.',
+    'Create a travel guide for a visit to Paris.',
+    'Write a poem about the changing seasons.',
+    'Explain the importance of clean energy sources.',
+    'List the essential camping gear for a wilderness adventure.',
+    'Write a short story about a time-traveling adventure.',
+    'Describe the history of the Renaissance.',
+    'Explain the principles of economics.',
+    'Create a business proposal for a new restaurant.',
+    'Write a letter to your future self 10 years from now.',
+    'List the ingredients for a classic lasagna.',
+    'Explain the concept of machine learning.',
+    'Describe the life and contributions of Martin Luther King Jr.',
+    'List the 10 most famous museums in the world.',
+    'Write a persuasive essay on the importance of environmental conservation.',
+    'Explain the process of geological erosion.',
+    'Create a menu for a vegan cafe.',
+    'Write a poem about the power of imagination.',
+    'Explain the significance of the Industrial Revolution.',
+    'List the items needed for a beach vacation.',
+    'Write a short biography of Charles Darwin.',
+    'Describe the impact of globalization on cultures.',
+    'Explain the principles of time management.',
+    'Create a playlist for a high-energy workout.',
+    'Write a blog post about the future of artificial intelligence.',
+    'Explain the process of DNA transcription.',
+    'List the 10 most iconic landmarks in Asia.',
+    'Write a book summary for a popular self-help book.',
+    'Describe the history of the ancient Greeks.',
+    'Explain the concept of social justice.',
+    'Create a recipe for a gourmet salad.',
+    'Write a screenplay for a romantic comedy movie.',
+    "Describe the layers of the Earth's atmosphere.",
+    'List the 10 most influential inventors in history.',
+    'Explain the process of plate tectonics.',
+    'Create a travel itinerary for a road trip across the USA.',
+    'Write a poem about the wonders of the natural world.',
+    'Explain the importance of sustainable agriculture.',
+    'List the essential hiking gear for a mountain expedition.',
+    'Write a short story about a futuristic dystopia.',
+    'Describe the history of the Middle Ages.',
+    'Write a letter to your future self, offering reflections on personal growth, achievements, and aspirations, as well as words of encouragement and guidance for your future journey.',
+    'List the ingredients for a classic chicken pot pie recipe, a beloved comfort food that combines tender chicken, vegetables, and a flaky pastry crust in a savory filling.',
+    'Explain the concept of artificial neural networks and their pivotal role in machine learning and artificial intelligence applications, from image recognition to natural language processing.',
+    'Describe the life and contributions of Albert Einstein, shedding light on his groundbreaking theories of relativity, his influence on modern physics, and his enduring legacy.',
+    "List the 10 most iconic landmarks in Australia, celebrating the country's breathtaking natural landscapes, architectural wonders, and unique cultural sites.",
+    'Write a persuasive speech on the importance of mental health awareness, advocating for destigmatization, access to mental health services, and compassionate support for those in need.',
+    "Explain the process of plate tectonics and its profound impact on the Earth's geology, including the formation of continents, mountain ranges, and the movement of tectonic plates.",
+    'Create a menu for a Mexican street food restaurant, featuring authentic and flavorful dishes that capture the essence of Mexican cuisine, from tacos to tamales.',
+    'Write a poem about the beauty of a moonlit night, exploring themes of serenity, reflection, and the enchanting ambiance created by the gentle glow of the moon.',
+    'Explain the significance of the Civil Rights Movement in the United States, highlighting the courageous individuals and leaders who paved the way for racial equality and justice.',
+    'List the items needed for a camping trip in the wilderness, ensuring that outdoor enthusiasts are well-equipped for their adventure in the great outdoors.',
+    'Write a short biography of Jane Austen, delving into her literary contributions and the enduring impact of her novels on literature, feminism, and societal norms.',
+    'Describe the impact of social media on political activism, discussing its role in mobilizing movements, shaping public discourse, and bringing about political change.',
+    'Explain the principles of effective leadership and management in the context of modern organizations, emphasizing the qualities and skills that define successful leaders.',
+    'Create a playlist for a high-energy workout session, curating motivating songs that provide the perfect soundtrack for a productive fitness routine.',
+    'Write a blog post about the future of artificial intelligence, exploring its potential applications in healthcare, transportation, and various industries that stand to be transformed.',
+    'Explain the process of DNA replication, shedding light on the intricate mechanism by which genetic information is faithfully duplicated during cell division.',
+    'List the 10 most famous landmarks in Asia, offering insights into their historical and cultural significance as well as their unique architectural features.',
+    'Write a book review for a contemporary novel that captivated your imagination, sharing your thoughts on the plot, characters, and overall impact of the book.',
+    'Describe the history of ancient Egypt, offering a glimpse into the rich civilization that thrived along the banks of the Nile River, from monumental architecture to hieroglyphic writing.',
+    'Explain the concept of cultural diversity, emphasizing the importance of embracing and celebrating diverse cultures and perspectives in a globalized world.',
+    'Create a recipe for a gourmet pasta dish that tantalizes the taste buds with a harmonious blend of flavors and textures, elevating a classic dish to culinary excellence.',
+    'Write a screenplay for a thrilling action movie set in a dystopian future, weaving together elements of suspense, adventure, and compelling characters in a post-apocalyptic world.',
+    "Describe the structure of the Earth's inner core, exploring its composition and the geophysical processes that drive the Earth's magnetic field and seismic activity.",
+    'List the 10 most significant technological innovations of the 21st century, from transformative advances in communication to breakthroughs in medical science and beyond.',
+    "Explain the process of geological earthquakes, delving into the geological forces and factors that contribute to the occurrence of seismic events and their impact on the Earth's surface.",
+    'Create a travel itinerary for a trip to Tokyo, Japan, offering recommendations for cultural experiences, sightseeing, and culinary delights in this vibrant metropolis.',
+    'Write a poem about the beauty of a sunset over the ocean, capturing the awe-inspiring colors and the sense of tranquility that descends as the sun dips below the horizon.',
+    'Explain the importance of conserving endangered species, highlighting the critical role these species play in maintaining ecosystem balance and the urgent need for conservation efforts.',
+    "List the essential gear for a camping and hiking adventure in a pristine natural wilderness, ensuring that outdoor enthusiasts are well-prepared for their journey into nature's beauty.",
+    'Write a short story about a group of explorers embarking on a thrilling quest to uncover a hidden treasure, navigating treacherous landscapes and solving ancient riddles along the way.',
+    'Describe the history of the Byzantine Empire, exploring its cultural achievements, architectural marvels such as the Hagia Sophia, and its enduring influence on Eastern Europe and beyond.',
+    'Explain the principles of ethical leadership, delving into the moral responsibilities of leaders in various domains and the impact of ethical leadership on organizations and society.',
+    'Create a business proposal for a sustainable fashion brand committed to eco-friendly practices, ethical sourcing, and transparency in the fashion industry.',
+    'Write a letter to your future self, offering reflections on personal growth, achievements, and aspirations, as well as words of encouragement and guidance for your future journey.',
+    'List the ingredients for a classic chicken pot pie recipe, a beloved comfort food that combines tender chicken, vegetables, and a flaky pastry crust in a savory filling.',
+    'Explain the concept of artificial neural networks and their pivotal role in machine learning and artificial intelligence applications, from image recognition to natural language processing.',
+    'Describe the life and contributions of Albert Einstein, shedding light on his groundbreaking theories of relativity, his influence on modern physics, and his enduring legacy.',
+    "List the 10 most iconic landmarks in Australia, celebrating the country's breathtaking natural landscapes, architectural wonders, and unique cultural sites.",
+    'Write a persuasive speech on the importance of mental health awareness, advocating for destigmatization, access to mental health services, and compassionate support for those in need.',
+    "Explain the process of plate tectonics and its profound impact on the Earth's geology, including the formation of continents, mountain ranges, and the movement of tectonic plates.",
+    'Create a menu for a Mexican street food restaurant, featuring authentic and flavorful dishes that capture the essence of Mexican cuisine, from tacos to tamales.',
+    'Write a poem about the beauty of a moonlit night, exploring themes of serenity, reflection, and the enchanting ambiance created by the gentle glow of the moon.',
+    'Explain the significance of the Civil Rights Movement in the United States, highlighting the courageous individuals and leaders who paved the way for racial equality and justice.',
+    'List the items needed for a camping trip in the wilderness, ensuring that outdoor enthusiasts are well-equipped for their adventure in the great outdoors.',
+    'Write a short biography of Jane Austen, delving into her literary contributions and the enduring impact of her novels on literature, feminism, and societal norms.',
+    'Describe the impact of social media on political activism, discussing its role in mobilizing movements, shaping public discourse, and bringing about political change.',
+    'Explain the principles of effective leadership and management in the context of modern organizations, emphasizing the qualities and skills that define successful leaders.',
+    'Create a playlist for a high-energy workout session, curating motivating songs that provide the perfect soundtrack for a productive fitness routine.',
+    'Write a blog post about the future of artificial intelligence, exploring its potential applications in healthcare, transportation, and various industries that stand to be transformed.',
+    'Explain the process of DNA replication, shedding light on the intricate mechanism by which genetic information is faithfully duplicated during cell division.',
+    'List the 10 most famous landmarks in Asia, offering insights into their historical and cultural significance as well as their unique architectural features.',
+    'Write a book review for a contemporary novel that captivated your imagination, sharing your thoughts on the plot, characters, and overall impact of the book.',
+    'Describe the history of ancient Egypt, offering a glimpse into the rich civilization that thrived along the banks of the Nile River, from monumental architecture to hieroglyphic writing.',
+    'Explain the concept of cultural diversity, emphasizing the importance of embracing and celebrating diverse cultures and perspectives in a globalized world.',
+    'Create a recipe for a gourmet pasta dish that tantalizes the taste buds with a harmonious blend of flavors and textures, elevating a classic dish to culinary excellence.',
+    'Write a screenplay for a thrilling action movie set in a dystopian future, weaving together elements of suspense, adventure, and compelling characters in a post-apocalyptic world.',
+    "Describe the structure of the Earth's inner core, exploring its composition and the geophysical processes that drive the Earth's magnetic field and seismic activity.",
+    'List the 10 most significant technological innovations of the 21st century, from transformative advances in communication to breakthroughs in medical science and beyond.',
+    "Explain the process of geological earthquakes, delving into the geological forces and factors that contribute to the occurrence of seismic events and their impact on the Earth's surface.",
+    'Create a travel itinerary for a trip to Tokyo, Japan, offering recommendations for cultural experiences, sightseeing, and culinary delights in this vibrant metropolis.',
+    'Write a poem about the beauty of a sunset over the ocean, capturing the awe-inspiring colors and the sense of tranquility that descends as the sun dips below the horizon.',
+    'Explain the importance of conserving endangered species, highlighting the critical role these species play in maintaining ecosystem balance and the urgent need for conservation efforts.',
+    "List the essential gear for a camping and hiking adventure in a pristine natural wilderness, ensuring that outdoor enthusiasts are well-prepared for their journey into nature's beauty.",
+    'Write a short story about a group of explorers embarking on a thrilling quest to uncover a hidden treasure, navigating treacherous landscapes and solving ancient riddles along the way.',
+    'Describe the history of the Byzantine Empire, exploring its cultural achievements, architectural marvels such as the Hagia Sophia, and its enduring influence on Eastern Europe and beyond.',
+    'Explain the principles of ethical leadership, delving into the moral responsibilities of leaders in various domains and the impact of ethical leadership on organizations and society.',
+    'Create a business proposal for a sustainable fashion brand committed to eco-friendly practices, ethical sourcing, and transparency in the fashion industry.',
+    'Write a letter to your future self, offering reflections on personal growth, achievements, and aspirations, as well as words of encouragement and guidance for your future journey.',
   ]
   async with aiohttp.ClientSession() as session:
-    await asyncio.gather(*[send_request(url, it, prompt, session, 'llama', max_new_tokens=2048) for it, prompt in enumerate(prompts)])
+    await asyncio.gather(*[send_request(url, it, prompt, session, args.model, max_new_tokens=2048) for it, prompt in enumerate(prompts)])
   return 0
 
+
 if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument('--generate', default=False, action='store_true', help='Whether to test with stream endpoint.')
-  args = parser.parse_args()
-  raise SystemExit(asyncio.run(main(args)))
+  parser.add_argument('--model', default='llama', choices=openllm.CONFIG_MAPPING_NAMES.keys(), action='store', help='Whether to test with stream endpoint.')
+  raise SystemExit(asyncio.run(main(parser.parse_args())))
diff --git a/cz.py b/cz.py
index ac604cfb..b16f449a 100755
--- a/cz.py
+++ b/cz.py
@@ -9,12 +9,14 @@ from tabulate import tabulate
 
 TOKEN_WHITELIST = [token.OP, token.NAME, token.NUMBER, token.STRING]
 
+
 def run_cz(dir: str, package: str):
   headers = ['Name', 'Lines', 'Tokens/Line']
   table = []
   for path, _, files in os.walk(os.path.join(dir, 'src', package)):
     for name in files:
-      if not name.endswith('.py'): continue
+      if not name.endswith('.py'):
+        continue
       filepath = os.path.join(path, name)
       with tokenize.open(filepath) as file_:
         tokens = [t for t in tokenize.generate_tokens(file_.readline) if t.type in TOKEN_WHITELIST]
@@ -23,15 +25,21 @@ def run_cz(dir: str, package: str):
   print(f'\n{"=" * 80}\n')
   print(tabulate([headers, *sorted(table, key=lambda x: -x[1])], headers='firstrow', floatfmt='.1f') + '\n')
   print(
-      tabulate([(dir_name, sum([x[1] for x in group])) for dir_name, group in itertools.groupby(sorted([(x[0].rsplit('/', 1)[0], x[1]) for x in table]), key=lambda x: x[0])],
-               headers=['Directory', 'LOC'],
-               floatfmt='.1f'))
+    tabulate(
+      [(dir_name, sum([x[1] for x in group])) for dir_name, group in itertools.groupby(sorted([(x[0].rsplit('/', 1)[0], x[1]) for x in table]), key=lambda x: x[0])],
+      headers=['Directory', 'LOC'],
+      floatfmt='.1f',
+    )
+  )
   print(f'total line count for {package}: {sum([x[1] for x in table])}\n')
 
+
 def main() -> int:
   run_cz('openllm-python', 'openllm')
   run_cz('openllm-core', 'openllm_core')
   run_cz('openllm-client', 'openllm_client')
   return 0
 
-if __name__ == '__main__': raise SystemExit(main())
+
+if __name__ == '__main__':
+  raise SystemExit(main())
diff --git a/examples/bentoml-demo/service.py b/examples/bentoml-demo/service.py
index 25cf3b4b..5548042a 100644
--- a/examples/bentoml-demo/service.py
+++ b/examples/bentoml-demo/service.py
@@ -7,6 +7,7 @@ llm = openllm.LLM('facebook/opt-2.7b')
 
 svc = bentoml.Service(name='llm-service', runners=[llm.runner])
 
+
 @svc.api(input=bentoml.io.Text(), output=bentoml.io.Text())
 async def prompt(input_text: str) -> str:
   generation = await llm.generate(input_text)
diff --git a/examples/langchain-chains-demo/service.py b/examples/langchain-chains-demo/service.py
index 8f1e48c6..79d7c116 100644
--- a/examples/langchain-chains-demo/service.py
+++ b/examples/langchain-chains-demo/service.py
@@ -9,21 +9,25 @@ from pydantic import BaseModel
 import bentoml
 from bentoml.io import JSON, Text
 
+
 class Query(BaseModel):
   industry: str
   product_name: str
   keywords: t.List[str]
   llm_config: t.Dict[str, t.Any]
 
+
 def gen_llm(model_name: str, model_id: str | None = None, **attrs: t.Any) -> OpenLLM:
   lc_llm = OpenLLM(model_name=model_name, model_id=model_id, embedded=False, **attrs)
   lc_llm.runner.download_model()
   return lc_llm
 
+
 llm = gen_llm('llama', model_id='TheBloke/Llama-2-13B-chat-GPTQ', quantize='gptq')
 
-prompt = PromptTemplate(input_variables=['industry', 'product_name', 'keywords'],
-                        template="""
+prompt = PromptTemplate(
+  input_variables=['industry', 'product_name', 'keywords'],
+  template="""
 You are a Facebook Ads Copywriter with a strong background in persuasive
 writing and marketing. You craft compelling copy that appeals to the target
 audience's emotions and needs, peruading them to take action or make a
@@ -35,15 +39,16 @@ Industry: {industry}
 Product: {product_name}
 Keywords: {keywords}
 Facebook Ads copy:
-    """)
+    """,
+)
 chain = LLMChain(llm=llm, prompt=prompt)
 
 svc = bentoml.Service('fb-ads-copy', runners=[llm.runner])
 
-SAMPLE_INPUT = Query(industry='SAAS',
-                     product_name='BentoML',
-                     keywords=['open source', 'developer tool', 'AI application platform', 'serverless', 'cost-efficient'],
-                     llm_config=llm.runner.config.model_dump())
+SAMPLE_INPUT = Query(
+  industry='SAAS', product_name='BentoML', keywords=['open source', 'developer tool', 'AI application platform', 'serverless', 'cost-efficient'], llm_config=llm.runner.config.model_dump()
+)
+
 
 @svc.api(input=JSON.from_sample(sample=SAMPLE_INPUT), output=Text())
 def generate(query: Query):
diff --git a/examples/langchain-tools-demo/service.py b/examples/langchain-tools-demo/service.py
index 6f38a6a9..4df4ab7e 100644
--- a/examples/langchain-tools-demo/service.py
+++ b/examples/langchain-tools-demo/service.py
@@ -13,6 +13,7 @@ tools = load_tools(['serpapi'], llm=llm)
 agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)
 svc = bentoml.Service('langchain-openllm', runners=[llm.runner])
 
+
 @svc.api(input=Text.from_sample(sample=SAMPLE_INPUT), output=Text())
 def chat(input_text: str):
   return agent.run(input_text)
diff --git a/examples/openai_client.py b/examples/openai_client.py
index 61fa548d..353b929b 100644
--- a/examples/openai_client.py
+++ b/examples/openai_client.py
@@ -1,39 +1,46 @@
-from __future__ import annotations
 import os
 
+# NOTE: Make sure to install openai>1
 import openai
-import importlib.util
-
-SUPPORT_LOGPROBS = str(os.getenv('ENABLE_LOGPROBS', default=importlib.util.find_spec('vllm') is not None)).upper() in ['TRUE', '1', 'YES', 'Y', 'ON']
-
-openai.api_base = os.getenv('OPENLLM_ENDPOINT', 'http://localhost:3000') + '/v1'
-openai.api_key = 'na'
+from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam, ChatCompletionAssistantMessageParam
 
 MODEL = os.getenv('MODEL', 'na')  # XXX: CHANGE THIS TO THE MODEL USED AT $OPENLLM_ENDPOINT
-
-print('Model metadata (/v1/models):\n')
-print(openai.Model.list())
-
-print('\nOne-shot completion (/v1/completions):\n')
-print(openai.Completion.create(model=MODEL, prompt='Write a tagline for an ice cream shop.', max_tokens=12))
-
-print('\nStreaming completion (/v1/completions):\n')
-for chunk in openai.Completion.create(model=MODEL, prompt='Say this is a test', max_tokens=128, temperature=0.8, stream=True, logprobs=2 if SUPPORT_LOGPROBS else None):
-  print(chunk.choices[0].text, flush=True, end='')
-
+CLIENT = openai.OpenAI(base_url=os.getenv('OPENLLM_ENDPOINT', 'http://localhost:3000') + '/v1', api_key='na')
 MESSAGES = [
-  {
-    'role': 'system',
-    'content': 'You are acting as Ernest Hemmingway. You must answers questions that assume the identity of the writer. You must always return factual information and will not tolerate misleading ideology.',
-  },
-  {'role': 'user', 'content': 'Hi there!'},
-  {'role': 'assistant', 'content': 'Yes?'},
-  {'role': 'user', 'content': 'What is the meaning of life?'},
+  ChatCompletionSystemMessageParam(
+    role='system',
+    content='You are acting as Ernest Hemmingway. You must answers questions that assume the identity of the writer. You must always return factual information and will not tolerate misleading ideology.',
+  ),
+  ChatCompletionUserMessageParam(role='user', content='Hi there!'),
+  ChatCompletionAssistantMessageParam(role='assistant', content='Yes?'),
+  ChatCompletionUserMessageParam(role='user', content='What is the meaning of life?'),
 ]
-ARGS = dict(max_tokens=512, temperature=0.83, top_p=0.72, top_k=12, logprobs=2 if SUPPORT_LOGPROBS else None)
 
-print('\nOne-shot chat completion (/v1/chat/completions):\n')
-print(openai.ChatCompletion.create(model=MODEL, messages=MESSAGES, **ARGS))
-print('\nStreaming chat completion (/v1/chat/completions):\n')
-for chunk in openai.ChatCompletion.create(model=MODEL, messages=MESSAGES, stream=True, **ARGS):
-  print(chunk.choices[0].delta.content, flush=True, end='')
+
+def completions() -> bool:
+  print(CLIENT.completions.create(model=MODEL, prompt='Write a tagline for an ice cream shop.', max_tokens=64, temperature=0.83, logprobs=1))
+  for chunk in CLIENT.completions.create(model=MODEL, prompt='Say this is a test', stream=True, max_tokens=64, temperature=0.83, logprobs=2):
+    print(chunk.choices[0].text, flush=True, end='')
+  return True
+
+
+def chat_completions() -> bool:
+  print(CLIENT.chat.completions.create(MESSAGES, MODEL, max_tokens=512, temperature=0.83, top_p=0.72))
+  for chunk in CLIENT.chat.completions.create(MESSAGES, MODEL, stream=True, max_tokens=512, temperature=0.83, top_p=0.72):
+    content = chunk.choices[0].delta.content
+    if content:
+      print(content, flush=True, end='')
+  return True
+
+
+def main() -> int:
+  print('Model: %s' % CLIENT.models.list().model_dump())
+  if not completions():
+    return 1
+  if not chat_completions():
+    return 1
+  return 0
+
+
+if __name__ == '__main__':
+  raise SystemExit(main())
diff --git a/openllm-core/src/openllm_core/_configuration.py b/openllm-core/src/openllm_core/_configuration.py
index 6bc6c272..4569712a 100644
--- a/openllm-core/src/openllm_core/_configuration.py
+++ b/openllm-core/src/openllm_core/_configuration.py
@@ -515,9 +515,9 @@ class ModelSettings(t.TypedDict, total=False):
   # the target generation_config class to be used.
   fine_tune_strategies: t.Tuple[t.Dict[str, t.Any], ...]
 
-  # tokenizer_class is the custom tokenizer class for this given LLM
-  tokenizer_class: t.Optional[str]
+  # Chat models related configuration
   conversation: t.Optional[t.Dict[str, t.Any]]
+  add_generation_prompt: bool
 
 _transformed_type: DictStrAny = {'fine_tune_strategies': t.Dict[AdapterType, FineTuneConfig], 'conversation': Conversation}
 
@@ -554,7 +554,8 @@ class _ModelSettingsAttr:
                       model_type='causal_lm',
                       trust_remote_code=False,
                       requirements=None,
-                      tokenizer_class=None,
+                      conversation=dict(system_message='', roles=('', ''), sep_style=SeparatorStyle.NO_COLON_SINGLE, sep=''),
+                      add_generation_prompt=False,
                       timeout=int(36e6),
                       service_name='',
                       workers_per_resource=1.)))
@@ -579,8 +580,8 @@ class _ModelSettingsAttr:
     timeout: int
     workers_per_resource: t.Union[int, float]
     fine_tune_strategies: t.Dict[AdapterType, FineTuneConfig]
-    tokenizer_class: t.Optional[str]
     conversation: Conversation
+    add_generation_prompt: bool
     # update-config-stubs.py: attrs stop
 
 def structure_settings(cl_: type[LLMConfig], cls: type[_ModelSettingsAttr]) -> _ModelSettingsAttr:
@@ -785,10 +786,10 @@ class _ConfigAttr:
         '''
     __openllm_fine_tune_strategies__: t.Dict[AdapterType, FineTuneConfig] = Field(None)
     '''The fine-tune strategies for this given LLM.'''
-    __openllm_tokenizer_class__: t.Optional[str] = Field(None)
-    '''Optional tokenizer class for this given LLM. See Llama for example.'''
     __openllm_conversation__: Conversation = Field(None)
     '''The conversation class for this given LLM to determine its chat templates.'''
+    __openllm_add_generation_prompt__: bool = Field(None)
+    '''Whether to add generation prompt token for formatting chat templates. This arguments will be used for chat-based models.'''
     # update-config-stubs.py: special stop
 
 class _ConfigBuilder:
@@ -1148,9 +1149,9 @@ class LLMConfig(_ConfigAttr):
   @overload
   def __getitem__(self, item: t.Literal['fine_tune_strategies']) -> t.Dict[AdapterType, FineTuneConfig]: ...
   @overload
-  def __getitem__(self, item: t.Literal['tokenizer_class']) -> t.Optional[str]: ...
-  @overload
   def __getitem__(self, item: t.Literal['conversation']) -> Conversation: ...
+  @overload
+  def __getitem__(self, item: t.Literal['add_generation_prompt']) -> bool: ...
   # NOTE: generation_class, sampling_class and extras arguments
   @overload
   def __getitem__(self, item: t.Literal['generation_class']) -> t.Type[openllm_core.GenerationConfig]: ...
@@ -1251,8 +1252,6 @@ class LLMConfig(_ConfigAttr):
   @overload
   def __getitem__(self, item: t.Literal['use_beam_search']) -> bool: ...
   @overload
-  def __getitem__(self, item: t.Literal['stop']) -> t.List[str]: ...
-  @overload
   def __getitem__(self, item: t.Literal['ignore_eos']) -> bool: ...
   @overload
   def __getitem__(self, item: t.Literal['logprobs']) -> int: ...
@@ -1264,6 +1263,8 @@ class LLMConfig(_ConfigAttr):
   @overload
   def __getitem__(self, item: t.Literal['prompt_tuning']) -> dict[str, t.Any]: ...
   @overload
+  def __getitem__(self, item: t.Literal['multitask_prompt_tuning']) -> dict[str, t.Any]: ...
+  @overload
   def __getitem__(self, item: t.Literal['p_tuning']) -> dict[str, t.Any]: ...
   @overload
   def __getitem__(self, item: t.Literal['prefix_tuning']) -> dict[str, t.Any]: ...
@@ -1275,6 +1276,10 @@ class LLMConfig(_ConfigAttr):
   def __getitem__(self, item: t.Literal['adaption_prompt']) -> dict[str, t.Any]: ...
   @overload
   def __getitem__(self, item: t.Literal['ia3']) -> dict[str, t.Any]: ...
+  @overload
+  def __getitem__(self, item: t.Literal['loha']) -> dict[str, t.Any]: ...
+  @overload
+  def __getitem__(self, item: t.Literal['lokr']) -> dict[str, t.Any]: ...
   # update-config-stubs.py: stop
   # fmt: on
 
diff --git a/openllm-core/src/openllm_core/_schemas.py b/openllm-core/src/openllm_core/_schemas.py
index 93753fd9..cc289c8c 100644
--- a/openllm-core/src/openllm_core/_schemas.py
+++ b/openllm-core/src/openllm_core/_schemas.py
@@ -136,15 +136,21 @@ class GenerationOutput:
                prompt_logprobs=None,
                request_id=gen_random_uuid())
 
+  @staticmethod
+  def _preprocess_sse_message(data: str) -> str:
+    proc = [line[6:] for line in data.strip().split('\n') if line.startswith('data: ')]
+    if not proc: return data
+    if len(proc) > 1: raise ValueError('Multiple data found in SSE message.')
+    return proc[0]
+
   @classmethod
-  def from_sse(cls, sse_message: str) -> GenerationOutput:
-    data = [line[6:] for line in sse_message.strip().split('\n') if line.startswith('data: ')]
-    if not data: raise ValueError('No data found in SSE message.')
-    if len(data) > 1: raise ValueError('Multiple data found in SSE message.')
+  def from_runner(cls, data: str) -> GenerationOutput:
+    data = cls._preprocess_sse_message(data)
+    if not data: raise ValueError('No data found from messages.')
     try:
-      return converter.structure(orjson.loads(data[0]), cls)
+      return converter.structure(orjson.loads(data), cls)
     except orjson.JSONDecodeError as e:
-      raise ValueError(f'Failed to parse JSON from SSE message: {sse_message!r}') from e
+      raise ValueError(f'Failed to parse JSON from SSE message: {data!r}') from e
 
   @classmethod
   def from_vllm(cls, request_output: vllm.RequestOutput) -> GenerationOutput:
diff --git a/openllm-core/src/openllm_core/config/configuration_llama.py b/openllm-core/src/openllm_core/config/configuration_llama.py
index 0a767e84..c600e206 100644
--- a/openllm-core/src/openllm_core/config/configuration_llama.py
+++ b/openllm-core/src/openllm_core/config/configuration_llama.py
@@ -76,7 +76,6 @@ class LlamaConfig(openllm_core.LLMConfig):
       'url': 'https://github.com/facebookresearch/llama',
       'architecture': 'LlamaForCausalLM',
       'requirements': ['fairscale', 'sentencepiece', 'scipy'],
-      'tokenizer_class': 'LlamaTokenizerFast',
       'default_id': 'NousResearch/llama-2-7b-hf',
       'serialisation': 'safetensors',
       # NOTE: see https://huggingface.co/blog/codellama#conversational-instructions
diff --git a/openllm-core/src/openllm_core/config/configuration_mistral.py b/openllm-core/src/openllm_core/config/configuration_mistral.py
index 3d65008e..6a1689b9 100644
--- a/openllm-core/src/openllm_core/config/configuration_mistral.py
+++ b/openllm-core/src/openllm_core/config/configuration_mistral.py
@@ -37,6 +37,7 @@ class MistralConfig(openllm_core.LLMConfig):
       'url': 'https://huggingface.co/docs/transformers/v4.35.0/en/model_doc/mistral#overview',
       'default_id': 'mistralai/Mistral-7B-Instruct-v0.1',
       'architecture': 'MistralForCausalLM',
+      'add_generation_prompt': True,
       'model_ids': ['mistralai/Mistral-7B-v0.1', 'mistralai/Mistral-7B-Instruct-v0.1', 'amazon/MistralLite', 'HuggingFaceH4/zephyr-7b-beta', 'HuggingFaceH4/zephyr-7b-alpha'],
   }
 
diff --git a/openllm-python/pyproject.toml b/openllm-python/pyproject.toml
index ce7377ea..2f4c4a66 100644
--- a/openllm-python/pyproject.toml
+++ b/openllm-python/pyproject.toml
@@ -112,7 +112,7 @@ gptq = ["auto-gptq[triton]>=0.4.2", "optimum>=1.12.0"]
 grpc = ["openllm-client[grpc]"]
 llama = ["fairscale", "sentencepiece", "scipy"]
 mpt = ["triton", "einops"]
-openai = ["openai[embeddings]", "tiktoken"]
+openai = ["openai[datalib]>=1", "tiktoken"]
 opt = ["transformers>=4.34.0"]
 playground = ["jupyter", "notebook", "ipython", "jupytext", "nbformat"]
 starcoder = ["bitsandbytes"]
diff --git a/openllm-python/src/openllm/_llm.py b/openllm-python/src/openllm/_llm.py
index d3fc52d1..6dc2b611 100644
--- a/openllm-python/src/openllm/_llm.py
+++ b/openllm-python/src/openllm/_llm.py
@@ -353,15 +353,13 @@ class LLM(t.Generic[M, T]):
     if request_id is None: request_id = openllm_core.utils.gen_random_uuid()
     previous_texts, previous_num_tokens = [''] * config['n'], [0] * config['n']
     async for out in self.runner.generate_iterator.async_stream(prompt_token_ids, request_id, stop, adapter_name, **config.model_dump(flatten=True)):
-      generated = GenerationOutput.from_sse(out).with_options(prompt=prompt)
+      generated = GenerationOutput.from_runner(out).with_options(prompt=prompt)
       delta_outputs = t.cast(t.List[CompletionChunk], [None] * len(generated.outputs))
       if generated.finished: break
       for output in generated.outputs:
         i = output.index
-        delta_tokens = output.token_ids[previous_num_tokens[i]:]
-        delta_text = output.text[len(previous_texts[i]):]
-        previous_texts[i] = output.text
-        previous_num_tokens[i] = len(output.token_ids)
+        delta_tokens, delta_text = output.token_ids[previous_num_tokens[i]:], output.text[len(previous_texts[i]):]
+        previous_texts[i], previous_num_tokens[i] = output.text, len(output.token_ids)
         delta_outputs[i] = output.with_options(text=delta_text, token_ids=delta_tokens)
       yield generated.with_options(outputs=delta_outputs)
 
diff --git a/openllm-python/src/openllm/_runners.py b/openllm-python/src/openllm/_runners.py
index 28c104ce..c12e6807 100644
--- a/openllm-python/src/openllm/_runners.py
+++ b/openllm-python/src/openllm/_runners.py
@@ -80,7 +80,7 @@ class vLLMRunnable(bentoml.Runnable):
     async for request_output in self.model.generate(None, sampling_params, request_id, prompt_token_ids):
       # XXX: Need to write a hook for serialisation None correctly
       if request_output.prompt_logprobs is not None: request_output.prompt_logprobs = [it if it else {} for it in request_output.prompt_logprobs]
-      yield f'data: {GenerationOutput.from_vllm(request_output).model_dump_json()}\n\n'
+      yield GenerationOutput.from_vllm(request_output).model_dump_json()
 
 class PyTorchRunnable(bentoml.Runnable):
   SUPPORTED_RESOURCES = ('nvidia.com/gpu', 'amd.com/gpu', 'cpu')
@@ -101,7 +101,7 @@ class PyTorchRunnable(bentoml.Runnable):
                               **attrs: t.Any) -> t.AsyncGenerator[str, None]:
     if adapter_name is not None: self.model.set_adapter(adapter_name)
     async for generation_output in self.forward(prompt_token_ids, request_id, stop=stop, **attrs):
-      yield f'data: {generation_output.model_dump_json()}\n\n'
+      yield generation_output.model_dump_json()
 
   async def forward(self, prompt_token_ids: list[int], request_id: str, stop: str | t.Iterable[str] | None = None, **attrs: t.Any) -> t.AsyncGenerator[GenerationOutput, None]:
     from ._generation import is_partial_stop
diff --git a/openllm-python/src/openllm/entrypoints/openai.py b/openllm-python/src/openllm/entrypoints/openai.py
index b559fbee..5ee40d28 100644
--- a/openllm-python/src/openllm/entrypoints/openai.py
+++ b/openllm-python/src/openllm/entrypoints/openai.py
@@ -128,7 +128,7 @@ async def create_chat_completions(req: Request, llm: openllm.LLM[M, T]) -> Respo
 
   model_name, request_id = request.model, gen_random_uuid('chatcmpl')
   created_time = int(time.monotonic())
-  prompt = llm.tokenizer.apply_chat_template(request.messages, tokenize=False)
+  prompt = llm.tokenizer.apply_chat_template(request.messages, tokenize=False, add_generation_prompt=llm.config['add_generation_prompt'])
   logger.debug('Prompt: %r', prompt)
   config = llm.config.with_openai_request(request)
 
diff --git a/openllm-python/src/openllm/playground/features.py b/openllm-python/src/openllm/playground/features.py
index dabbe822..575c24ab 100644
--- a/openllm-python/src/openllm/playground/features.py
+++ b/openllm-python/src/openllm/playground/features.py
@@ -40,8 +40,3 @@ async def main() -> int:
 def _mp_fn(index: t.Any):  # type: ignore
   # For xla_spawn (TPUs)
   asyncio.run(main())
-
-if openllm.utils.in_notebook():
-  await main()
-else:
-  raise SystemExit(asyncio.run(main()))
diff --git a/tools/dependencies.py b/tools/dependencies.py
index 2840af4e..8e44ce86 100755
--- a/tools/dependencies.py
+++ b/tools/dependencies.py
@@ -135,7 +135,7 @@ FINE_TUNE_DEPS = ['peft>=0.6.0', 'bitsandbytes', 'datasets', 'accelerate', 'trl'
 FLAN_T5_DEPS = ['transformers>=4.34.0']
 OPT_DEPS = ['transformers>=4.34.0']
 GRPC_DEPS = ['openllm-client[grpc]']
-OPENAI_DEPS = ['openai[embeddings]', 'tiktoken']
+OPENAI_DEPS = ['openai[datalib]>=1', 'tiktoken']
 AGENTS_DEPS = ['transformers[agents]>=4.34.0', 'diffusers', 'soundfile']
 PLAYGROUND_DEPS = ['jupyter', 'notebook', 'ipython', 'jupytext', 'nbformat']
 GGML_DEPS = ['ctransformers']
diff --git a/tools/update-config-stubs.py b/tools/update-config-stubs.py
index 357171ff..30970e70 100755
--- a/tools/update-config-stubs.py
+++ b/tools/update-config-stubs.py
@@ -53,6 +53,7 @@ _value_docstring = {
             openllm start gpt-neox --model-id stabilityai/stablelm-tuned-alpha-3b
             ```''',
     'conversation': '''The conversation class for this given LLM to determine its chat templates.''',
+    'add_generation_prompt': '''Whether to add generation prompt token for formatting chat templates. This arguments will be used for chat-based models.''',
     'backend': '''List of supported backend for this given LLM class. Currently, we support "pt" and "vllm".''',
     'serialisation': '''Default serialisation format for different models. Some will default to use the legacy 'bin'. ''',
     'url': 'The resolved url for this LLMConfig.',
@@ -77,7 +78,6 @@ _value_docstring = {
         By default, it is set to 1.
         ''',
     'fine_tune_strategies': 'The fine-tune strategies for this given LLM.',
-    'tokenizer_class': 'Optional tokenizer class for this given LLM. See Llama for example.',
 }
 
 _transformed = {'fine_tune_strategies': 't.Dict[AdapterType, FineTuneConfig]', 'conversation': 'Conversation'}