Update app.py

app.py

@@ -3,71 +3,59 @@ import torch
 from threading import Thread
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-# Model and Tokenizer Configuration 
+# Model and tokenizer config
 MODEL_REPO_ID = "LGAI-EXAONE/EXAONE-4.0-1.2B"
 
 print("✅ Starting application...")
 
-# Load the model with bfloat16 to save memory
 try:
     print(f"🔄 Loading tokenizer from '{MODEL_REPO_ID}'...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO_ID)
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_REPO_ID,
+        trust_remote_code=True  
+    )
     print("✅ Tokenizer loaded successfully.")
 
     print(f"🔄 Loading model '{MODEL_REPO_ID}' with torch_dtype=torch.bfloat16...")
-    
+
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_REPO_ID,
         torch_dtype=torch.bfloat16,
-        device_map="auto" 
+        device_map="auto", 
+        trust_remote_code=True  
     )
     print("✅ Model loaded successfully.")
 
 except Exception as e:
     print(f"❌ Error loading model or tokenizer: {e}")
-    # Exit if model fails to load, as the app is unusable.
     raise
 
-# Streaming Chat Function
+# Streaming chat function 
 def user_input_handler(user_message, history):
-    """Handles user input by appending it to the history."""
     return "", history + [[user_message, None]]
 
 def bot_stream(history):
-    """
-    Generates the bot's response using a streaming approach.
-    This function runs the model in a separate thread to avoid blocking the UI.
-    """
-    print(f"📝 History received: {history}")
-    # The last message is the user's prompt.
     user_message = history[-1][0]
-    history[-1][1] = "" # Initialize the bot's response field.
+    history[-1][1] = "" # Initialize the bot's response field
 
-    # Format the conversation history into the model's expected chat format.
+    # Format the conversation history into the model's expected chat format
     messages = []
-    for human, assistant in history[:-1]: # All but the last interaction
+    for human, assistant in history[:-1]:
         messages.append({"role": "user", "content": human})
-        if assistant: # Assistant message might be None
+        if assistant:
             messages.append({"role": "assistant", "content": assistant})
     messages.append({"role": "user", "content": user_message})
 
-    try:
-        # Apply the chat template to format the prompt.
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    except Exception as e:
-        print(f"⚠️ Warning: Could not apply chat template. Using basic formatting. Error: {e}")
-        # Fallback for models without a registered chat template
-        prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages]) + "\nassistant:"
-
-    print("➡️ Generated Prompt for Model:\n" + prompt)
-
-    # Tokenize the formatted prompt.
+    # Apply the chat template to format the prompt correctly
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    # Tokenize the formatted prompt
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    # Use TextIteratorStreamer for non-blocking, token-by-token generation.
+    # Use TextIteratorStreamer for non-blocking, token-by-token generation
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-    # Set up the generation parameters in a dictionary.
+    # Set up the generation parameters in a dictionary
     generation_kwargs = dict(
         inputs,
         streamer=streamer,
@@ -78,19 +66,19 @@ def bot_stream(history):
         eos_token_id=tokenizer.eos_token_id,
     )
 
-    # Run the generation in a separate thread to avoid blocking the Gradio UI.
+    # Run the generation in a separate thread
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    # Yield each new token to the Gradio chat interface as it's generated.
+    # Yield each new token to the Gradio chat interface as it's generated
     for token in streamer:
         history[-1][1] += token
         yield history
 
-# Gradio User Interface
+# Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css="footer {display: none !important}") as demo:
     gr.Markdown("## 🤖 EXAONE-4.0-1.2B")
-    gr.Markdown("This demo runs the standard `LGAI-EXAONE/EXAONE-4.0-1.2B` model using the `transformers` library.")
+    gr.Markdown("This demo runs `LGAI-EXAONE/EXAONE-4.0-1.2B` model")
 
     chatbot = gr.Chatbot(label="Chat History", height=600, bubble_full_width=False)
     with gr.Row():
@@ -104,7 +92,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), cs
 
     clear_btn = gr.ClearButton([msg, chatbot], value="🗑️ Clear Chat")
 
-    # Event Handlers
+    # Event Handlers - when a message is submitted 
     msg.submit(user_input_handler, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot_stream, chatbot, chatbot
     )
@@ -113,4 +101,4 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), cs
     )
 
 demo.queue()
-demo.launch(debug=True) 
+demo.launch(debug=True)