Update app.py

app.py

@@ -4,7 +4,7 @@ import logging
 import os
 from functools import lru_cache
 from threading import Thread
-from typing import Generator, List, Tuple
+from typing import Generator
 
 import gradio as gr
 import regex
@@ -41,7 +41,7 @@ def load_model_and_tokenizer():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
-        torch_dtype=torch.float16,
+        dtype=torch.bfloat16,
         device_map="auto",
     )
     logger.info("Model loaded successfully")
@@ -54,24 +54,21 @@ DEFAULT_LANG = "he"
 
 labels = {
     "he": {
-        "title": "מתמטיבוט 🧮",
+        "title": '<h1 dir="rtl" style="text-align: right;">מתמטיבוט 🧮</h1>',
         "intro": (
-            """
-ברוכים הבאים לדמו! 💡 כאן תוכלו להתרשם **ממודל השפה החדש** שלנו; מודל בגודל 4 מיליארד פרמטרים שאומן לענות על שאלות מתמטיות בעברית, על המחשב שלכם, ללא חיבור לרשת.
-
-קישור למודל, פרטים נוספים, יצירת קשר ותנאי שימוש:
-
-https://huggingface.co/Intel/hebrew-math-tutor-v1
-
------
-"""
+            '<div dir="rtl" style="text-align: right;">\n\n'
+            "ברוכים הבאים לדמו! 💡 כאן תוכלו להתרשם **ממודל השפה החדש** שלנו; מודל בגודל 4 מיליארד פרמטרים שאומן לענות על שאלות מתמטיות בעברית, על המחשב שלכם, ללא חיבור לרשת.\n\n"
+            "קישור למודל, פרטים נוספים, יצירת קשר ותנאי שימוש:\n\n"
+            "https://huggingface.co/Intel/hebrew-math-tutor-v1\n\n"
+            "-----\n\n"
+            "</div>"
         ),
         "select_label": "בחרו שאלה מוכנה או צרו שאלה חדשה:",
         "new_question": "שאלה חדשה...",
         "text_label": "שאלה:",
         "placeholder": "הזינו את השאלה כאן...",
         "send": "שלח",
-        "reset": "שיחה חדשה",
+        "reset": "נקה",
         "toggle_to": "English 🇬🇧",
         "predefined": [
             "שאלה חדשה...",
@@ -83,10 +80,10 @@ https://huggingface.co/Intel/hebrew-math-tutor-v1
         "thinking_prefix": "🤔 חושב",
         "thinking_done": "🤔 *תהליך החשיבה הושלם, מכין תשובה...*",
         "final_label": "📝 תשובה סופית:",
-        "chat_label": "צ'אט",
+        "answer_label": "תשובה:",
     },
     "en": {
-        "title": "MathBot 🧮",
+        "title": "<h1>MathBot 🧮</h1>",
         "intro": (
             """
 Welcome to the demo! 💡 Here you can try our **new language model** — a 4-billion-parameter model trained to answer math questions in Hebrew while maintaining its English capabilities. It runs locally on your machine without requiring an internet connection.
@@ -103,7 +100,7 @@ https://huggingface.co/Intel/hebrew-math-tutor-v1
         "text_label": "Question:",
         "placeholder": "Type your question here...",
         "send": "Send",
-        "reset": "New Conversation",
+        "reset": "Clear",
         "toggle_to": "עברית 🇮🇱",
         "predefined": [
             "New question...",
@@ -115,12 +112,12 @@ https://huggingface.co/Intel/hebrew-math-tutor-v1
         "thinking_prefix": "🤔 Thinking",
         "thinking_done": "🤔 *Thinking complete, preparing answer...*",
         "final_label": "📝 Final answer:",
-        "chat_label": "Chat",
+        "answer_label": "Answer:",
     },
 }
 
 
-def dir_and_alignment(lang: str) -> Tuple[str, str]:
+def dir_and_alignment(lang: str) -> tuple[str, str]:
     if lang == "he":
         return "rtl", "right"
     return "ltr", "left"
@@ -175,7 +172,7 @@ def build_assistant_markdown(
     direction, align = dir_and_alignment(lang)
     localized = labels[lang]
 
-    parts: List[str] = []
+    parts = []
     if thinking_text:
         details = _details_template.format(
             dir=direction,
@@ -193,57 +190,52 @@ def build_assistant_markdown(
     return "\n\n".join(parts)
 
 
-@spaces.GPU
+# @spaces.GPU
 def handle_user_message(
     user_input: str,
     lang: str,
-    chat_history: List[Tuple[str, str]] | None,
 ) -> Generator[tuple, None, None]:
     lang = lang or DEFAULT_LANG
     localized = labels[lang]
-    chat_history = chat_history or []
     prompt = (user_input or "").strip()
     if not prompt:
         yield (
-            chat_history,
+            "",
             "",
             localized["new_question"],
-            chat_history,
         )
         return
 
-    formatted_user = wrap_text_with_direction(prompt, lang)
-    chat_history = chat_history + [(formatted_user, "")]
     dropdown_reset = localized["new_question"]
-    yield chat_history, "", dropdown_reset, chat_history
+    yield "", "", dropdown_reset
 
     system_prompt = build_system_prompt(lang)
 
-    # Format as chat template
+    # Apply chat template with just user message
     chat_messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt},
     ]
 
-    # Apply chat template
     input_text = tokenizer.apply_chat_template(
         chat_messages,
-        tokenize=False,
+        tokenize=True,
         add_generation_prompt=True,
+        enable_thinking=True,
     )
 
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    inputs = torch.tensor([input_text]).to(model.device)
 
-    thinking_buffer = ""
+    thinking_buffer = "<think>"
     thinking_text: str | None = None
     final_answer = ""
     response_fallback = ""
-    in_thinking = False
+    in_thinking = True
 
     try:
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         generation_kwargs = dict(
-            **inputs,
+            inputs=inputs,
             streamer=streamer,
             max_new_tokens=2400,
             temperature=0.6,
@@ -261,8 +253,8 @@ def handle_user_message(
 
             response_fallback += delta
 
-            if "<think>" in delta:
-                in_thinking = True
+            # if "<think>" in delta:
+            #     in_thinking = True
 
             if in_thinking:
                 thinking_buffer += delta
@@ -280,34 +272,30 @@ def handle_user_message(
                     thinking_text=thinking_text,
                 )
 
-            chat_history[-1] = (formatted_user, current_answer)
-            yield chat_history, "", dropdown_reset, chat_history
+            yield current_answer, "", dropdown_reset
 
         thread.join()
 
     except Exception as exc:
         error_html = wrap_text_with_direction(f"⚠️ Error generating response: {exc}", lang)
-        chat_history[-1] = (formatted_user, error_html)
-        yield chat_history, prompt, dropdown_reset, chat_history
+        yield error_html, prompt, dropdown_reset
         return
 
     if not final_answer:
         final_answer = response_fallback
 
-    chat_history[-1] = (
-        formatted_user,
-        build_assistant_markdown(lang=lang, final_answer=final_answer, thinking_text=thinking_text),
+    final_markdown = build_assistant_markdown(
+        lang=lang, final_answer=final_answer, thinking_text=thinking_text
     )
-    yield chat_history, "", dropdown_reset, chat_history
+    yield final_markdown, "", dropdown_reset
 
 
-def reset_conversation(lang: str):
+def clear_answer(lang: str):
     localized = labels[lang]
     return (
-        [],
+        "",
         "",
         localized["new_question"],
-        [],
     )
 
 
@@ -322,62 +310,140 @@ def toggle_language_and_rebuild(lang: str):
     """Toggle language and return new values for all UI components"""
     new_lang = "en" if lang == "he" else "he"
     localized = labels[new_lang]
+    # direction, align = dir_and_alignment(new_lang)
+    dropdown_class = "rtl" if new_lang == "he" else "ltr"
+    textbox_class = "rtl" if new_lang == "he" else "ltr"
     return (
-        new_lang,  # lang_state
-        f"# {localized['title']}",  # title_md
-        localized["intro"],  # intro_md
-        gr.Dropdown(
-            label=localized["select_label"],
+        new_lang,
+        localized["title"],
+        localized["intro"],
+        wrap_text_with_direction(localized["select_label"], new_lang, emphasized=True),
+        gr.update(
             choices=localized["predefined"],
             value=localized["new_question"],
-            interactive=True,
-        ),  # preset_dropdown
-        gr.Textbox(
-            label=localized["text_label"],
-            placeholder=localized["placeholder"],
-            lines=5,
-            value="",
-        ),  # question_box
-        localized["send"],  # send_button value
-        localized["reset"],  # reset_button value
-        localized["toggle_to"],  # lang_button value
-        localized["chat_label"],  # chatbot label
+            elem_classes=[dropdown_class],
+        ),
+        wrap_text_with_direction(localized["text_label"], new_lang, emphasized=True),
+        gr.update(
+            placeholder=localized["placeholder"], value="", elem_classes=[textbox_class]
+        ),
+        localized["send"],
+        localized["reset"],
+        wrap_text_with_direction(localized["answer_label"], new_lang, emphasized=True),
+        localized["toggle_to"],
     )
 
 
+CUSTOM_CSS = """
+body {
+    font-family: 'Rubik', 'Segoe UI', 'Helvetica Neue', Arial, sans-serif;
+}
+details > summary {
+    cursor: pointer;
+}
+.gradio-container .prose p {
+    margin-bottom: 0.5rem;
+}
+[dir="rtl"] {
+    direction: rtl;
+    text-align: right;
+}
+[dir="ltr"] {
+    direction: ltr;
+    text-align: left;
+}
+/* Alignment helpers for inputs and dropdowns using elem_classes */
+.rtl textarea, .rtl .gr-dropdown {
+    direction: rtl !important;
+    text-align: right !important;
+}
+.ltr textarea, .ltr .gr-dropdown {
+    direction: ltr !important;
+    text-align: left !important;
+}
+.rtl input, .rtl select, .rtl .gr-dropdown, .rtl .gradio-dropdown {
+    direction: rtl !important;
+    text-align: right !important;
+}
+.ltr input, .ltr select, .ltr .gr-dropdown, .ltr .gradio-dropdown {
+    direction: ltr !important;
+    text-align: left !important;
+}
+
+/* Visual container for the answer area */
+.answer-box {
+    background: #ffffff !important;
+    border: 1px solid rgba(0,0,0,0.08);
+    border-radius: 10px;
+    padding: 12px 16px;
+    box-shadow: none;
+    margin-top: 8px;
+}
+.answer-box .prose {
+    margin: 0;
+}
+"""
+
+
 def build_demo() -> gr.Blocks:
     localized = labels[DEFAULT_LANG]
     with gr.Blocks(title="Hebrew Math Tutor") as demo:
+        # Inject custom CSS
+        gr.HTML(f"<style>{CUSTOM_CSS}</style>")
+
         lang_state = gr.State(DEFAULT_LANG)
-        chat_state = gr.State([])
 
-        title_md = gr.Markdown(f"# {localized['title']}")
+        title_md = gr.Markdown(localized["title"])
         intro_md = gr.Markdown(localized["intro"])
 
-        with gr.Row():
-            preset_dropdown = gr.Dropdown(
-                label=localized["select_label"],
-                choices=localized["predefined"],
-                value=localized["new_question"],
-                interactive=True,
+        select_label_md = gr.Markdown(
+            wrap_text_with_direction(
+                localized["select_label"], DEFAULT_LANG, emphasized=True
             )
-            lang_button = gr.Button(localized["toggle_to"], variant="secondary")
+        )
+
+        preset_dropdown = gr.Dropdown(
+            label="",
+            choices=localized["predefined"],
+            value=localized["new_question"],
+            interactive=True,
+            show_label=False,
+            elem_classes=("rtl" if DEFAULT_LANG == "he" else "ltr"),
+        )
+
+        question_label_md = gr.Markdown(
+            wrap_text_with_direction(
+                localized["text_label"], DEFAULT_LANG, emphasized=True
+            )
+        )
 
         question_box = gr.Textbox(
-            label=localized["text_label"],
+            label="",
             placeholder=localized["placeholder"],
             lines=5,
+            show_label=False,
+            elem_classes=("rtl" if DEFAULT_LANG == "he" else "ltr"),
         )
 
         with gr.Row():
             reset_button = gr.Button(localized["reset"], variant="secondary")
             send_button = gr.Button(localized["send"], variant="primary")
+            lang_button = gr.Button(localized["toggle_to"], variant="secondary")
 
-        chatbot = gr.Chatbot(
-            label=localized["chat_label"],
-            height=520,
-            render_markdown=True,
+        answer_label_md = gr.Markdown(
+            wrap_text_with_direction(
+                localized["answer_label"], DEFAULT_LANG, emphasized=True
+            )
         )
+        with gr.Group(elem_classes="answer-box"):
+            answer_box = gr.Markdown(
+                label="",
+                show_label=False,
+                latex_delimiters=[
+                    {"left": "$$", "right": "$$", "display": True},
+                    {"left": "$", "right": "$", "display": False},
+                ],
+            )
 
         preset_dropdown.change(
             fn=sync_question_text,
@@ -386,30 +452,40 @@ def build_demo() -> gr.Blocks:
         )
 
         reset_button.click(
-            fn=reset_conversation,
+            fn=clear_answer,
             inputs=[lang_state],
-            outputs=[chatbot, question_box, preset_dropdown, chat_state],
+            outputs=[answer_box, question_box, preset_dropdown],
         )
 
         send_button.click(
             fn=handle_user_message,
-            inputs=[question_box, lang_state, chat_state],
-            outputs=[chatbot, question_box, preset_dropdown, chat_state],
+            inputs=[question_box, lang_state],
+            outputs=[answer_box, question_box, preset_dropdown],
         )
 
         question_box.submit(
             fn=handle_user_message,
-            inputs=[question_box, lang_state, chat_state],
-            outputs=[chatbot, question_box, preset_dropdown, chat_state],
+            inputs=[question_box, lang_state],
+            outputs=[answer_box, question_box, preset_dropdown],
         )
 
-        # Simplified language toggle - just updates the state
-        # In a production app, you'd want to rebuild the entire UI
-        # For now, we'll just update the state (labels won't change dynamically)
+        # Language toggle - updates state and UI labels
         lang_button.click(
-            fn=lambda lang: "en" if lang == "he" else "he",
+            fn=toggle_language_and_rebuild,
             inputs=[lang_state],
-            outputs=[lang_state],
+            outputs=[
+                lang_state,
+                title_md,
+                intro_md,
+                select_label_md,
+                preset_dropdown,
+                question_label_md,
+                question_box,
+                send_button,
+                reset_button,
+                answer_label_md,
+                lang_button,
+            ],
         )
 
     return demo
@@ -419,21 +495,3 @@ demo = build_demo()
 
 if __name__ == "__main__":
     demo.queue().launch()
-
-
-demo = build_demo()
-
-CUSTOM_CSS = """
-body {
-    font-family: 'Rubik', 'Segoe UI', 'Helvetica Neue', Arial, sans-serif;
-}
-details > summary {
-    cursor: pointer;
-}
-.gradio-container .prose p {
-    margin-bottom: 0.5rem;
-}
-"""
-
-if __name__ == "__main__":
-    demo.queue().launch(css=CUSTOM_CSS)