up

README.md

@@ -1,8 +1,8 @@
 ---
-title: Vigogne-Chat
+title: Vigogne Chat
 emoji: 🦙
-colorFrom: purple
-colorTo: yellow
+colorFrom: yellow
+colorTo: purple
 sdk: gradio
 sdk_version: 3.27.0
 app_file: app.py

app.py

@@ -13,6 +13,8 @@ python vigogne/demo/demo_chat.py \
     --lora_model_name_or_path bofenghuang/vigogne-chat-7b
 """
 
+import json
+
 # import datetime
 import logging
 import os
@@ -20,10 +22,6 @@ import re
 from threading import Event, Thread
 from typing import List, Optional
 
-
-# from uuid import uuid4
-
-import json
 import gradio as gr
 
 # import requests
@@ -33,13 +31,16 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
+    LlamaTokenizer,
     StoppingCriteriaList,
     TextIteratorStreamer,
 )
-
 from vigogne.constants import ASSISTANT, USER
-from vigogne.preprocess import generate_inference_chat_prompt
 from vigogne.inference.inference_utils import StopWordsCriteria
+from vigogne.preprocess import generate_inference_chat_prompt
+
+# from uuid import uuid4
+
 
 logging.basicConfig(
     format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
@@ -97,7 +98,9 @@ def main(
     server_port: Optional[str] = None,
     share: bool = False,
 ):
-    tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path, padding_side="right", use_fast=False)
+    # tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path, padding_side="right", use_fast=False)
+    tokenizer_class = LlamaTokenizer if "llama" in base_model_name_or_path else AutoTokenizer
+    tokenizer = tokenizer_class.from_pretrained(base_model_name_or_path, padding_side="right", use_fast=False)
 
     if device == "cuda":
         model = AutoModelForCausalLM.from_pretrained(
@@ -302,12 +305,7 @@ def main(
                 elem_classes=["disclaimer"],
             )
 
-        submit_event = msg.submit(
-            fn=user,
-            inputs=[msg, chatbot],
-            outputs=[msg, chatbot],
-            queue=False,
-        ).then(
+        submit_event = msg.submit(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False,).then(
             fn=bot,
             inputs=[
                 chatbot,
@@ -321,12 +319,7 @@ def main(
             outputs=chatbot,
             queue=True,
         )
-        submit_click_event = submit.click(
-            fn=user,
-            inputs=[msg, chatbot],
-            outputs=[msg, chatbot],
-            queue=False,
-        ).then(
+        submit_click_event = submit.click(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False,).then(
             fn=bot,
             inputs=[
                 chatbot,
@@ -352,4 +345,5 @@ def main(
     demo.queue(max_size=128, concurrency_count=2)
     demo.launch(enable_queue=True, share=share, server_name=server_name, server_port=server_port)
 
-main()
+
+main(base_model_name_or_path="decapoda-research/llama-7b-hf", lora_model_name_or_path="bofenghuang/vigogne-chat-7b")