Update app.py

app.py

@@ -1,28 +1,29 @@
 import os
-from threading import Thread
 from typing import Iterator
 import gradio as gr
 import torch
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer
+from openai import OpenAI
+import json
+import uuid
 
 
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+EXAONE_TOKEN = os.environ.get("EXAONE_TOKEN", None)
+EXAONE_2_4B = os.environ.get("EXAONE_2_4B", None)
+EXAONE_7_8B = os.environ.get("EXAONE_7_8B", None)
+EXAONE_32B = os.environ.get("EXAONE_32B", None)
 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>
 
 ##### <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://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://arxiv.org/abs/2412.04862>technical report</a></center>
-
-#### <center> EXAONE-3.5-32B-Instruct Demo Coming Soon.. </center>
 """
 
 
@@ -31,7 +32,10 @@ EXAMPLES = [
     ["스스로를 자랑해 봐"],
 ]
 BOT_AVATAR = "EXAONE_logo.png"
-selected_model = gr.Radio(value="https://jps6tfdq34ydttbh.us-east4.gcp.endpoints.huggingface.cloud",visible=False)
+selected_model = gr.Radio(value=["2.4B", EXAONE_2_4B],visible=False)
+id_ = {"id": str(uuid.uuid4())}
+model_history = {"model_history": []}
+
 ADDITIONAL_INPUTS = [
     gr.Textbox(
             value="You are EXAONE model from LG AI Research, a helpful assistant.",
@@ -59,18 +63,12 @@ ADDITIONAL_INPUTS = [
         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")
 
+
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -78,13 +76,10 @@ def generate(
     max_new_tokens: int = 512,
     temperature: float = 0.6,
     top_p: float = 0.9,
-    top_k: int = 50,
-    selected_model: str = "https://jps6tfdq34ydttbh.us-east4.gcp.endpoints.huggingface.cloud",
+    selected_model: list = ["2.4b", EXAONE_2_4B],
 ) -> Iterator[str]:
-    print(f'model: {selected_model}')
     messages = [{"role":"system","content": system_prompt}]
-    print(f'message: {message}')
-    print(f'chat_history: {chat_history}')
+    
     for user, assistant in chat_history:
         messages.extend(
             [
@@ -94,30 +89,34 @@ def generate(
         )
     messages.append({"role": "user", "content": message})
 
+    if not chat_history:
+        id_['id'] = str(uuid.uuid4())
+        model_history["model_history"] = []    
+    model_history["model_history"].append(selected_model[0])
+    
     input_ids = tokenizer.apply_chat_template(
         messages, 
         add_generation_prompt=True, 
         return_tensors="pt"
     )
-    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.")
-    
-    messages = tokenizer.decode(input_ids[0])
-    
-    client = InferenceClient(selected_model, token=HF_TOKEN)
-
-    gen_kwargs = dict(
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        stop=["[|endofturn|]"]
-    )
 
-    output = client.text_generation(messages, **gen_kwargs)
+    client = OpenAI(api_key=EXAONE_TOKEN, base_url="https://api.friendli.ai/dedicated/v1")
+    response = client.chat.completions.create(
+            messages=messages, 
+            model=selected_model[1],
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stream=True,
+    )
+    outputs = ''
+    for r in response:
+        token = r.choices[0].delta.content
+        if token is not None:
+            outputs += token
+        yield outputs
 
-    return output
+    print(json.dumps({"id": id_['id'], "messages": messages, "output": outputs, "model": model_history}, ensure_ascii=False))
 
 
 def radio1_change(model_size):
@@ -126,8 +125,9 @@ def radio1_change(model_size):
 
 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
+        "2.4B": ["2.4B", EXAONE_2_4B],
+        "7.8B": ["7.8B", EXAONE_7_8B],
+        "32B": ["32B", EXAONE_32B],
     }
     return endpoint_url_dict[model_size]
 
@@ -144,16 +144,16 @@ chat_interface = gr.ChatInterface(
         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("""<p align="center"><img src="https://huggingface.co/spaces/LGAI-EXAONE/EXAONE-3.5-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"]
+        model_size = ["2.4B", "7.8B", "32B"]
         radio1 = gr.Radio(choices=model_size, label="EXAONE-3.5-Instruct", value=model_size[0])
     
     radio1.change(radio1_change, inputs=radio1, outputs=markdown)