Update app.py

app.py

@@ -1,5 +1,3 @@
-# app.py: Fine-tune internally + Gradio UI
-
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
 from datasets import Dataset
 import gradio as gr
@@ -36,7 +34,8 @@ model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 # ------------------------------
 def preprocess(examples):
     model_inputs = tokenizer(examples['input'], truncation=True, padding='max_length', max_length=64)
-    labels = tokenizer(examples['output'], truncation=True, padding='max_length', max_length=64)
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(examples['output'], truncation=True, padding='max_length', max_length=64)
     model_inputs['labels'] = labels['input_ids']
     return model_inputs
 
@@ -65,11 +64,12 @@ trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=tokenized_dataset,
-    tokenizer=None,  # avoid FutureWarning
+    tokenizer=tokenizer,  # FIXED
     data_collator=data_collator
 )
 
-if not os.path.exists(OUTPUT_DIR):
+# Fine-tune only if directory is missing or empty
+if not os.path.exists(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR):
     trainer.train()
     model.save_pretrained(OUTPUT_DIR)
     tokenizer.save_pretrained(OUTPUT_DIR)
@@ -93,13 +93,14 @@ def respond(user_input, chat_history):
     return chat_history, chat_history
 
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 style='text-align:center'>💊 Health Remedies Chatbot</h1>")
+    gr.Markdown("<h1 style='text-align:center; color:#4CAF50;'>💊 Health Remedies Chatbot</h1>")
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type your message here...")
-    state = gr.State()
+    state = gr.State([])
     with gr.Row():
-        send = gr.Button("Send")
+        msg = gr.Textbox(placeholder="Type your message here...", scale=8)
+        send = gr.Button("Send", scale=2)
     send.click(respond, [msg, state], [chatbot, state])
     msg.submit(respond, [msg, state], [chatbot, state])
 
 demo.launch()
+