Update app.py

app.py

@@ -5,41 +5,73 @@ import gradio as gr
 import torch
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from huggingface_hub import InferenceClient
+
 
-MODEL_LIST = ["LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL = MODEL_LIST[0]
+MODEL = "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct"
+MAX_NEW_TOKENS = 4096
+DEFAULT_MAX_NEW_TOKENS = 512
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "16384"))
+
 
 DESCRIPTION = """\
 # <center> EXAONE 3.5: Series of Large Language Models for Real-world Use Cases </center>
 
-##### We hope EXAONE continues to advance Expert AI with its effectiveness and bilingual skills.
-
-<center>This is an official demo of <a href=https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct>LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct</a>, fine-tuned for instruction following.</center>
+##### <center> We hope EXAONE continues to advance Expert AI with its effectiveness and bilingual skills. </center>
 
-<center>👋 For more details, please check <a href=https://www.lgresearch.ai/blog/view?seq=507>our blog</a> or <a href=https://arxiv.org/abs/2412.04862>technical report</a></center>
+<center>👋 For more details, please check <a href=https://huggingface.co/collections/LGAI-EXAONE/exaone-35-674d0e1bb3dcd2ab6f39dbb4>EXAONE-3.5 collections</a>, <a href=https://www.lgresearch.ai/blog/view?seq=507>our blog</a> or <a href=https://www.lgresearch.ai/data/upload/tech_report/ko/Technical_report_EXAONE_3.5.pdf>technical report</a></center>
 
-#### <center> EXAONE-3.5-2.4B-Instruct and EXAONE-3.5-32B-Instruct Demo Coming Soon.. </center>
+#### <center> EXAONE-3.5-32B-Instruct Demo Coming Soon.. </center>
 """
 
-MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 512
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "16384"))
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL,
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-    device_map="auto",
-)
 
-model.eval()
+EXAMPLES = [
+    ["Explain how wonderful you are"],
+    ["스스로를 자랑해 봐"],
+]
+BOT_AVATAR = "EXAONE_logo.png"
+selected_model = gr.Radio(value="https://jps6tfdq34ydttbh.us-east4.gcp.endpoints.huggingface.cloud",visible=False)
+ADDITIONAL_INPUTS = [
+    gr.Textbox(
+            value="You are EXAONE model from LG AI Research, a helpful assistant.",
+            label="System Prompt",
+            render=False,
+    ),
+    gr.Slider(
+        label="Max new tokens",
+        minimum=1,
+        maximum=MAX_NEW_TOKENS,
+        step=1,
+        value=DEFAULT_MAX_NEW_TOKENS,
+    ),
+    gr.Slider(
+        label="Temperature",
+        minimum=0.1,
+        maximum=2.0,
+        step=0.1,
+        value=0.7,
+    ),
+    gr.Slider(
+        label="Top-p (nucleus sampling)",
+        minimum=0.05,
+        maximum=1.0,
+        step=0.05,
+        value=0.9,
+    ),
+    gr.Slider(
+        label="Top-k",
+        minimum=1,
+        maximum=1000,
+        step=1,
+        value=1,
+    ),
+    selected_model
+]
+
+tokenizer = AutoTokenizer.from_pretrained("LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct")
 
 
-@spaces.GPU()
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -48,7 +80,9 @@ def generate(
     temperature: float = 0.6,
     top_p: float = 0.9,
     top_k: int = 50,
+    selected_model: str = "https://jps6tfdq34ydttbh.us-east4.gcp.endpoints.huggingface.cloud",
 ) -> Iterator[str]:
+    print(f'model: {selected_model}')
     messages = [{"role":"system","content": system_prompt}]
     print(f'message: {message}')
     print(f'chat_history: {chat_history}')
@@ -69,88 +103,64 @@ def generate(
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
         gr.Warning(f"Trimmed input from messages as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    input_ids = input_ids.to(model.device)
+    
+    messages = tokenizer.decode(input_ids[0])
+    
+    client = InferenceClient(selected_model, token=HF_TOKEN)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
+    gen_kwargs = dict(
         max_new_tokens=max_new_tokens,
-        do_sample=False if top_k == 1 else True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
-        num_beams=1,
-        repetition_penalty=1.0,
+        stop=["[|endofturn|]"]
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    output = client.text_generation(messages, **gen_kwargs)
 
+    return output
 
-BOT_AVATAR = "EXAONE_logo.png"
 
-chatbot = gr.Chatbot(
-    label="EXAONE-3.5-7.8B-Instruct",
-    avatar_images=[None, BOT_AVATAR],
-    layout="bubble",
-    bubble_full_width=False
-)
+def radio1_change(model_size):
+    return f"<center><font size=5>EXAONE-3.5-{model_size}-instruct</center>"
+
+
+def choices_model(model_size):
+    endpoint_url_dict = {
+        "2.4B": "https://jps6tfdq34ydttbh.us-east4.gcp.endpoints.huggingface.cloud", # L4
+        "7.8B": "https://wafz6im0d595g715.us-east-1.aws.endpoints.huggingface.cloud", # L40S
+    }
+    return endpoint_url_dict[model_size]
+
 
 chat_interface = gr.ChatInterface(
-    fn=generate,
-    chatbot=chatbot,
-    additional_inputs=[
-        gr.Textbox(
-                value="You are EXAONE model from LG AI Research, a helpful assistant.",
-                label="System Prompt",
-                render=False,
-            ),
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=2.0,
-            step=0.1,
-            value=0.7,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
+        fn=generate,
+        chatbot=gr.Chatbot(
+            label="EXAONE-3.5-Instruct",
+            avatar_images=[None, BOT_AVATAR],
+            layout="bubble",
+            bubble_full_width=False
         ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=1,
-        ),
-    ],
-    stop_btn=None,
-    examples=[
-        ["Explain who you are"],
-        ["너의 소원을 말해봐"],
-    ],
-    cache_examples=False,
-)
-
-with gr.Blocks(css="style.css", fill_height=True) as demo:
-    gr.Markdown("""<p align="center"><img src="https://huggingface.co/spaces/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Demo/resolve/main/EXAONE_Symbol%2BBI_3d.png" style="margin-right: 20px; height: 60px"/><p>""")
+        additional_inputs=ADDITIONAL_INPUTS,
+        stop_btn=None,
+        examples=EXAMPLES,
+        cache_examples=False,
+    )
+
+
+with gr.Blocks(fill_height=True) as demo:
+    gr.Markdown("""<p align="center"><img src="https://huggingface.co/spaces/LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct-Demo/resolve/main/EXAONE_Symbol%2BBI_3d.png" style="margin-right: 20px; height: 50px"/><p>""")
     gr.Markdown(DESCRIPTION)
+
+    markdown = gr.Markdown("<center><font size=5>EXAONE-3.5-2.4B-instruct</center>")
+    with gr.Row():
+        model_size = ["2.4B", "7.8B"]
+        radio1 = gr.Radio(choices=model_size, label="EXAONE-3.5-Instruct", value=model_size[0])
+    
+    radio1.change(radio1_change, inputs=radio1, outputs=markdown)
+    radio1.change(choices_model, inputs=radio1, outputs=selected_model)
     chat_interface.render()
 
+    
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=25).launch()