Update app.py

app.py

@@ -9,17 +9,19 @@ hf_token = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 
 torch.set_num_threads(4)
 
-# Global model & tokenizer
+# Globals
 tokenizer = None
 model = None
+current_model_name = None
 
 # Load selected model
 def load_model(model_name):
-    global tokenizer, model
+    global tokenizer, model, current_model_name
     full_model_name = f"MaxLSB/{model_name}"
     tokenizer = AutoTokenizer.from_pretrained(full_model_name, token=hf_token)
     model = AutoModelForCausalLM.from_pretrained(full_model_name, token=hf_token)
     model.eval()
+    current_model_name = model_name
 
 # Initialize default model
 load_model("LeCarnet-8M")
@@ -49,7 +51,8 @@ def respond(message, max_tokens, temperature, top_p):
     response = ""
     for new_text in streamer:
         response += new_text
-        yield response
+        # prepend model name on its own line
+        yield f"**Model: {current_model_name}**\n\n{response}"
 
 # User input handler
 def user(message, chat_history):
@@ -69,8 +72,6 @@ def update_model(model_name):
     load_model(model_name)
     return []
 
-image_path = "static/le-carnet.png"
-
 # Gradio UI
 with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
     with gr.Row():
@@ -78,11 +79,9 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
         <div style="text-align: center; width: 100%;">
             <h1 style="margin: 0;">LeCarnet Demo 📊</h1>
         </div>
-        """)
+        """ )
 
-    # Main layout
     with gr.Row():
-        # Left column: Options
         with gr.Column(scale=1, min_width=150):
             model_selector = gr.Dropdown(
                 choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
@@ -94,7 +93,6 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
             top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p Sampling")
             clear_button = gr.Button("Clear Chat")
 
-        # Right column: Chat
         with gr.Column(scale=4):
             chatbot = gr.Chatbot(
                 bubble_full_width=False,
@@ -114,7 +112,6 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
                 label="Example Prompts"
             )
 
-    # Event handlers
     model_selector.change(fn=update_model, inputs=[model_selector], outputs=[])
     msg_input.submit(fn=user, inputs=[msg_input, chatbot], outputs=[msg_input, chatbot], queue=False).then(
         fn=bot, inputs=[chatbot, max_tokens, temperature, top_p], outputs=[chatbot]
@@ -122,4 +119,4 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
     clear_button.click(fn=lambda: None, inputs=None, outputs=chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.queue(default_concurrency_limit=10, max_size=10).launch(ssr_mode=False, max_threads=10)
+    demo.queue(default_concurrency_limit=10, max_size=10).launch(ssr_mode=False, max_threads=10)