update UI

app.py

@@ -5,9 +5,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 model_to_use = "thesven/Llama3-8B-SFT-code_bagel-bnb-4bit"
 
+# Initialize global variables for the tokenizer and model
+tokenizer = None
+model = None
+
 @spaces.GPU
-def start(n):
-    model_name_or_path = "thesven/Llama3-8B-SFT-code_bagel-bnb-4bit"
+def start():
+    global tokenizer, model
+    model_name_or_path = model_to_use
 
     # BitsAndBytesConfig for loading the model in 4-bit precision
     bnb_config = BitsAndBytesConfig(
@@ -23,14 +28,29 @@ def start(n):
         trust_remote_code=True,
         quantization_config=bnb_config
     )
-    model.pad_token = model.config.eos_token_id
+    model.pad_token_id = model.config.eos_token_id
+
+    return "Model loaded and ready!"
 
-    # Example response generation
-    input_text = "Hello, how are you?"
+def send_message(input_text):
+    global tokenizer, model
     input_ids = tokenizer(input_text, return_tensors='pt').input_ids.cuda()
     output = model.generate(inputs=input_ids, max_new_tokens=50)
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
     return generated_text
 
-demo = gr.Interface(fn=start, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# Chat with the Model")
+    
+    start_button = gr.Button("Start Model")
+    status_text = gr.Textbox(label="Status")
+    
+    start_button.click(start, inputs=None, outputs=status_text)
+    
+    chatbox = gr.Chatbot()
+    message = gr.Textbox(label="Your Message")
+    send_button = gr.Button("Send")
+    
+    send_button.click(send_message, inputs=message, outputs=chatbox)
+    
+demo.launch()