Update app.py

app.py

@@ -15,10 +15,8 @@ alpaca_prompt = """පහත දැක්වෙන්නේ යම් කාර
 ### ප්‍රතිචාරය:
 {}"""
 
-@spaces.GPU
-def infer(message, history, enable_history=False, max_new_tokens=1024):
+def load_model():
     global model, tokenizer
-    
     if model is None:
         tokenizer = AutoTokenizer.from_pretrained("manthilaffs/Gamunu-4B-Instruct-Alpha")
         model = AutoModelForCausalLM.from_pretrained(
@@ -27,6 +25,12 @@ def infer(message, history, enable_history=False, max_new_tokens=1024):
             device_map="auto",
         )
         model.eval()
+
+@spaces.GPU
+def generate_response(message, history, enable_history=False, max_new_tokens=1024):
+    global model, tokenizer
+    
+    load_model()
     
     # Add history only if enabled
     if enable_history and history:
@@ -77,6 +81,9 @@ def infer(message, history, enable_history=False, max_new_tokens=1024):
         
         if response_started:
             yield partial_text
+    
+    # Make sure thread completes
+    thread.join()
 
 # Custom CSS for styling with copy button
 custom_css = """
@@ -296,25 +303,36 @@ splash_html = """
 with gr.Blocks(css=custom_css, head=copy_js) as demo:
     gr.HTML(splash_html)
     
-    chat = gr.ChatInterface(
-        fn=lambda message, history: infer(message, history, enable_history.value, max_new_tokens.value),
-        title="🧠 Gamunu 4B Instruct - Demo",
-        theme=gr.themes.Default(text_size="sm"),
-        examples=[
-            ["හෙලෝ ගැමුණු! මම සමන්, ඔයාට කොහොමද?"],
-            ["ෆොටෝසින්තසිස් ක්‍රියාවලිය පැහැදිලි කරන්න."],
-            ["මෙම වාක්‍යය සිංහලයට පරිවර්තනය කරන්න: 'The sun rises in the east.'"],
-            ["'completed' තත්ත්වයේ ඇති වාර්තා ගණන ගණනය කිරීමට දත්ත සමුදා විමසුමක් (database query) ගොඩනඟන්න."],
-            ["ඔබ ගුරුවරයෙකු ලෙස ක්‍රියාකරන්න. ශිෂ්‍යයාට ඉතිහාසය උගන්වන්න."],
-            ["පහත ප්‍රකාශය ප්‍රංශ භාෂාවට පරිවර්තනය කරන්න. Laughter is the best medicine."],
-            ["ඝන වස්තුවක හා ද්‍රවයක පරිමාවන්හි වෙනස පැහැදිලි කරන්න."],
-            ["වෙස් මුහුණු කලාවේ ප්‍රධාන අංග මොනවාද? වර්තමානයේ මෙම කලාව ප්‍රචලිතව පවතින ප්‍රදේශ මොනවාද?"]
-        ]
-    )
+    enable_history = gr.State(value=False)
+    max_new_tokens = gr.State(value=512)
+    
+    with gr.Row():
+        with gr.Column():
+            chat = gr.ChatInterface(
+                fn=generate_response,
+                title="🧠 Gamunu 4B Instruct - Demo",
+                theme=gr.themes.Default(text_size="sm"),
+                type="messages",  # Use new messages format
+                examples=[
+                    ["හෙලෝ ගැමුණු! මම සමන්, ඔයාට කොහොමද?"],
+                    ["ෆොටෝසින්තසිස් ක්‍රියාවලිය පැහැදිලි කරන්න."],
+                    ["මෙම වාක්‍යය සිංහලයට පරිවර්තනය කරන්න: 'The sun rises in the east.'"],
+                    ["'completed' තත්ත්වයේ ඇති වාර්තා ගණන ගණනය කිරීමට දත්ත සමුදා විමසුමක් (database query) ගොඩනඟන්න."],
+                    ["ඔබ ගුරුවරයෙකු ලෙස ක්‍රියාකරන්න. ශිෂ්‍යයාට ඉතිහාසය උගන්වන්න."],
+                    ["පහත ප්‍රකාශය ප්‍රංශ භාෂාවට පරිවර්තනය කරන්න. Laughter is the best medicine."],
+                    ["ඝන වස්තුවක හා ද්‍රවයක පරිමාවන්හි වෙනස පැහැදිලි කරන්න."],
+                    ["වෙස් මුහුණු කලාවේ ප්‍රධාන අංග මොනවාද? වර්තමානයේ මෙම කලාව ප්‍රචලිතව පවතින ප්‍රදේශ මොනවාද?"]
+                ],
+                additional_inputs=[enable_history, max_new_tokens]
+            )
     
     with gr.Accordion("⚙️ Advanced Settings", open=False):
-        enable_history = gr.Checkbox(label="Enable chat history", value=False)
-        max_new_tokens = gr.Slider(64, 1024, value=512, step=32, label="🔢 Max New Tokens")
+        history_checkbox = gr.Checkbox(label="Enable chat history", value=False)
+        tokens_slider = gr.Slider(64, 1024, value=512, step=32, label="🔢 Max New Tokens")
+        
+        # Update state when controls change
+        history_checkbox.change(fn=lambda x: x, inputs=history_checkbox, outputs=enable_history)
+        tokens_slider.change(fn=lambda x: x, inputs=tokens_slider, outputs=max_new_tokens)
     
     gr.Markdown("""
 ---