Update app.py

app.py

@@ -1,15 +1,26 @@
-# Full script: Fine-tune + Gradio UI
+# app.py: Fine-tune internally + Gradio UI
 
-from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
+from datasets import Dataset
 import gradio as gr
 import torch
 import os
 
 # ------------------------------
-# 1. Dataset
+# 1. Internal Dataset
 # ------------------------------
-dataset = load_dataset('csv', data_files={'train': 'remedies.csv'}, split='train')
+data = [
+    {"input": "What should I do for a cold?", "output": "I'm sorry you're feeling unwell. Drink warm water, get rest, and consider vitamin C."},
+    {"input": "What to do if I have a headache?", "output": "I understand headaches are frustrating. Try meditation, rest, and stay hydrated."},
+    {"input": "My child has fever, what do I do?", "output": "Give paracetamol, keep them hydrated, and if fever persists, consult a doctor."},
+    {"input": "Who can I contact for fever treatment?", "output": "You can reach Dr. Ankit Verma at +91-9876543210 or Dr. Priya Singh at +91-9123456780."},
+    {"input": "I feel dizzy, what should I do?", "output": "Sit down, drink water, and rest. If it continues, see a doctor."},
+    {"input": "I am anxious and need help.", "output": "Feeling anxious is okay. Try deep breathing. You can also speak with Dr. Richa Nair at +91-9874455667."},
+    {"input": "I have mild back pain.", "output": "Gentle stretching and rest can help. For consultation, Dr. Amit Khanna +91-9988774455 is available."},
+    {"input": "My child has a cough and cold.", "output": "Dr. Sneha Kapoor at +91-9871122334 and Dr. Arjun Mehta at +91-9112233445 can assist. Keep your child warm."}
+]
+
+dataset = Dataset.from_list(data)
 
 # ------------------------------
 # 2. Model & Tokenizer
@@ -24,10 +35,8 @@ model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 # 3. Preprocess
 # ------------------------------
 def preprocess(examples):
-    inputs = examples['input']
-    targets = examples['output']
-    model_inputs = tokenizer(inputs, truncation=True, padding='max_length', max_length=64)
-    labels = tokenizer(targets, truncation=True, padding='max_length', max_length=64)
+    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)
     model_inputs['labels'] = labels['input_ids']
     return model_inputs
 
@@ -45,47 +54,52 @@ training_args = TrainingArguments(
     output_dir=OUTPUT_DIR,
     per_device_train_batch_size=4,
     num_train_epochs=3,
-    logging_steps=10,
+    logging_steps=5,
     save_steps=50,
     save_total_limit=2,
     fp16=torch.cuda.is_available(),
-    report_to="none",
+    report_to="none"
 )
 
 trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=tokenized_dataset,
-    tokenizer=tokenizer,
+    tokenizer=None,  # avoid FutureWarning
     data_collator=data_collator
 )
 
-# Fine-tune only if model is not already saved
 if not os.path.exists(OUTPUT_DIR):
     trainer.train()
-    trainer.save_model(OUTPUT_DIR)
+    model.save_pretrained(OUTPUT_DIR)
+    tokenizer.save_pretrained(OUTPUT_DIR)
 
 # ------------------------------
-# 6. Load fine-tuned model for inference
+# 6. Load fine-tuned model
 # ------------------------------
 model = AutoModelForSeq2SeqLM.from_pretrained(OUTPUT_DIR)
 tokenizer = AutoTokenizer.from_pretrained(OUTPUT_DIR)
 
-def respond(user_input):
+# ------------------------------
+# 7. Gradio UI (chat-style)
+# ------------------------------
+def respond(user_input, chat_history):
     inputs = tokenizer(user_input, return_tensors="pt")
     outputs = model.generate(**inputs, max_length=64)
     reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return reply
+    chat_history = chat_history or []
+    chat_history.append(("You", user_input))
+    chat_history.append(("Bot", reply))
+    return chat_history, chat_history
 
-# ------------------------------
-# 7. Gradio UI
-# ------------------------------
-ui = gr.Interface(
-    fn=respond,
-    inputs="text",
-    outputs="text",
-    title="Health Remedies Chatbot",
-    description="Ask health-related questions and get remedy suggestions!"
-)
+with gr.Blocks() as demo:
+    gr.Markdown("<h1 style='text-align:center'>💊 Health Remedies Chatbot</h1>")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Type your message here...")
+    state = gr.State()
+    with gr.Row():
+        send = gr.Button("Send")
+    send.click(respond, [msg, state], [chatbot, state])
+    msg.submit(respond, [msg, state], [chatbot, state])
 
-ui.launch()
+demo.launch()