Update app.py

app.py

@@ -2,17 +2,14 @@ import gradio as gr
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
 
-# Load dataset
 print("Loading dataset...")
 
 dataset = load_dataset("nazlicanto/persona-based-chat")
 
-# Choose a base model (DialoGPT)
 model_name = "microsoft/DialoGPT-medium"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Ensure pad_token is set
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
@@ -20,10 +17,8 @@ def preprocess_data(batch):
     inputs = ["\n".join(dialogue) + "\nBot: " + reference for dialogue, reference in zip(batch["dialogue"], batch["reference"])]
     return tokenizer(inputs, truncation=True, padding="max_length", max_length=128)
 
-# Apply preprocessing
 tokenized_dataset = dataset.map(preprocess_data, batched=True, remove_columns=dataset["train"].column_names)
 
-# Use validation if available; otherwise, split the train dataset
 if "validation" in tokenized_dataset:
     train_dataset = tokenized_dataset["train"]
     eval_dataset = tokenized_dataset["validation"]
@@ -32,7 +27,6 @@ else:
     train_dataset = train_test_split["train"]
     eval_dataset = train_test_split["test"]
 
-# Training arguments
 training_args = TrainingArguments(
     output_dir="./results",
     evaluation_strategy="steps",
@@ -49,7 +43,6 @@ training_args = TrainingArguments(
 )
 
 
-# Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -58,7 +51,6 @@ trainer = Trainer(
     tokenizer=tokenizer
 )
 
-# Train model
 def train_model():
     print("Starting training...")
     trainer.train()
@@ -66,7 +58,6 @@ def train_model():
     tokenizer.save_pretrained("trained_chatbot")
     return "Training Complete!"
 
-# Chatbot interface
 def chatbot(user_input):
     input_text = f"User: {user_input}\nBot:"
     inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
@@ -84,7 +75,6 @@ def chatbot(user_input):
     response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
     return response.strip()
 
-# Gradio UI
 iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", live=True)
 
 if __name__ == "__main__":