Add `zephyr`

app.py

@@ -2,6 +2,7 @@ import os
 
 import chainlit as cl
 import openai
+import torch
 from chainlit.input_widget import Select, Slider, Switch
 from langchain.chat_models import ChatOpenAI
 from llama_index import (
@@ -13,7 +14,13 @@ from llama_index import (
     load_index_from_storage,
 )
 from llama_index.callbacks.base import CallbackManager
+from llama_index.chat_engine import CondenseQuestionChatEngine
+from llama_index.embeddings import HuggingFaceEmbedding
 from llama_index.llms import ChatMessage, HuggingFaceLLM, MessageRole, OpenAI
+from llama_index.prompts import PromptTemplate
+from llama_index.query_engine import SubQuestionQueryEngine
+from llama_index.tools import QueryEngineTool, ToolMetadata
+from transformers import BitsAndBytesConfig
 
 
 def get_api_key():
@@ -66,7 +73,7 @@ async def start():
             Select(
                 id="Model",
                 label="Model",
-                values=["gpt-3.5-turbo", "gpt-4"],
+                values=["gpt-3.5-turbo", "gpt-4", "zephyr"],
                 initial_index=1,
             ),
             Slider(
@@ -86,10 +93,55 @@ async def start():
 async def setup_query_engine(settings):
     print("on_settings_update", settings)
 
-    llm = OpenAI(model=settings["Model"], temperature=settings["Temperature"])
-
+    def messages_to_prompt(messages):
+        prompt = ""
+        for message in messages:
+            if message.role == "system":
+                prompt += f"<|system|>\n{message.content}</s>\n"
+            elif message.role == "user":
+                prompt += f"<|user|>\n{message.content}</s>\n"
+            elif message.role == "assistant":
+                prompt += f"<|assistant|>\n{message.content}</s>\n"
+        if not prompt.startswith("<|system|>\n"):
+            prompt = "<|system|>\n</s>\n" + prompt
+        prompt = prompt + "<|assistant|>\n"
+        return prompt
+
+    if settings["Model"] == "zephyr":
+        model_name = "HuggingFaceH4/zephyr-7b-beta"
+        query_wrapper_prompt = PromptTemplate(
+            "<|system|>\n</s>\n<|user|>\n{query_str}</s>\n<|assistant|>\n"
+        )
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+        )
+        llm = HuggingFaceLLM(
+            model_name=model_name,
+            tokenizer_name=model_name,
+            query_wrapper_prompt=query_wrapper_prompt,
+            context_window=3900,
+            max_new_tokens=256,
+            model_kwargs={"quantization_config": quantization_config},
+            generate_kwargs={
+                "do_sample": True,
+                "temperature": settings["Temperature"],
+                "top_k": 50,
+                "top_p": 0.95,
+            },
+            messages_to_prompt=messages_to_prompt,
+            device_map="auto",
+        )
+    else:
+        llm = OpenAI(model=settings["Model"], temperature=settings["Temperature"])
+
+    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5")
     service_context = ServiceContext.from_defaults(
-        llm=llm, callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()])
+        llm=llm,
+        embed_model=embed_model,
+        callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()]),
     )
 
     query_engine = index.as_query_engine(

requirements.txt

@@ -2,3 +2,8 @@ chainlit
 llama-index
 trafilatura
 openai
+torch
+transformers
+accelerate
+scipy
+bitsandbytes