Upload app.py

app.py

@@ -0,0 +1,169 @@
+import os
+import time
+import gradio as gr
+import pandas as pd
+import torch
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Dataset
+from transformers import (
+    T5ForConditionalGeneration,
+    T5TokenizerFast,
+    DataCollatorForSeq2Seq,
+    Trainer,
+    TrainingArguments,
+    pipeline
+)
+
+DATA_PATH = "data/train.csv"
+DEFAULT_INPUT_COL = "text"
+DEFAULT_TARGET_COL = "label"
+
+class CSVDataset(Dataset):
+    def __init__(self, df, tokenizer, input_col, target_col, max_input_len=512, max_target_len=128, prefix="summarize: "):
+        self.inputs = df[input_col].astype(str).tolist()
+        self.targets = df[target_col].astype(str).tolist()
+        self.tokenizer = tokenizer
+        self.max_input_len = max_input_len
+        self.max_target_len = max_target_len
+        self.prefix = prefix
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, idx):
+        src = self.prefix + self.inputs[idx]
+        tgt = self.targets[idx]
+        model_inputs = self.tokenizer(
+            src, max_length=self.max_input_len, truncation=True, padding=False, return_tensors=None
+        )
+        with self.tokenizer.as_target_tokenizer():
+            labels = self.tokenizer(
+                tgt, max_length=self.max_target_len, truncation=True, padding=False, return_tensors=None
+            )
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+
+def run_training(base_model, epochs, batch_size, lr, warmup_steps, weight_decay, max_input_len, max_target_len, input_col, target_col, eval_ratio, grad_accum, fp16):
+    log_lines = []
+    def log(msg):
+        log_lines.append(msg)
+
+    if not os.path.exists(DATA_PATH):
+        return "data/train.csv not found.", ""
+
+    try:
+        df = pd.read_csv(DATA_PATH)
+    except Exception as e:
+        return f"Failed reading CSV: {e}", ""
+
+    for c in [input_col, target_col]:
+        if c not in df.columns:
+            return f"Column '{c}' not in CSV. Found: {list(df.columns)}", ""
+
+    log("Loading tokenizer & model...")
+    tok = T5TokenizerFast.from_pretrained(base_model)
+    mdl = T5ForConditionalGeneration.from_pretrained(base_model)
+
+    train_df, val_df = train_test_split(df, test_size=float(eval_ratio), random_state=42)
+    train_ds = CSVDataset(train_df, tok, input_col, target_col, max_input_len, max_target_len)
+    val_ds = CSVDataset(val_df, tok, input_col, target_col, max_input_len, max_target_len)
+
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tok, model=mdl)
+    output_dir = "checkpoint"
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=int(epochs),
+        per_device_train_batch_size=int(batch_size),
+        per_device_eval_batch_size=int(batch_size),
+        learning_rate=float(lr),
+        weight_decay=float(weight_decay),
+        warmup_steps=int(warmup_steps),
+        predict_with_generate=True,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        logging_steps=10,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        gradient_accumulation_steps=int(grad_accum),
+        fp16=bool(fp16),
+        report_to=[],
+    )
+
+    trainer = Trainer(
+        model=mdl,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=val_ds,
+        tokenizer=tok,
+        data_collator=data_collator
+    )
+
+    log("Starting training...")
+    trainer.train()
+
+    log("Saving model...")
+    trainer.save_model(output_dir)
+    tok.save_pretrained(output_dir)
+
+    return "\n".join(log_lines), "Training complete. Model saved to ./checkpoint"
+
+def make_pipe_from_checkpoint():
+    if not os.path.exists("checkpoint"):
+        raise RuntimeError("No checkpoint found. Train first.")
+    return pipeline("text2text-generation", model="checkpoint")
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🔧 Train & Share: Summarizer (FLAN‑T5)")
+    with gr.Tab("Train"):
+        gr.Markdown("Use defaults and click **Start Training**. This runs inside the Space.")
+        base_model = gr.Dropdown(choices=["google/flan-t5-small","google/flan-t5-base"], value="google/flan-t5-small", label="Base model")
+        epochs = gr.Slider(1, 6, value=2, step=1, label="Epochs")
+        batch_size = gr.Slider(2, 16, value=8, step=1, label="Batch size")
+        lr = gr.Textbox(value="5e-5", label="Learning rate")
+        warmup = gr.Textbox(value="100", label="Warmup steps")
+        wd = gr.Textbox(value="0.01", label="Weight decay")
+        max_in = gr.Slider(128, 1024, value=512, step=32, label="Max input length")
+        max_out = gr.Slider(32, 256, value=128, step=8, label="Max target length")
+        in_col = gr.Textbox(value=DEFAULT_INPUT_COL, label="Input column")
+        out_col = gr.Textbox(value=DEFAULT_TARGET_COL, label="Target column")
+        eval_ratio = gr.Textbox(value="0.1", label="Eval ratio (0-1)")
+        grad_accum = gr.Slider(1, 8, value=1, step=1, label="Gradient accumulation")
+        use_fp16 = gr.Checkbox(value=True, label="Use fp16 (GPU only)")
+        train_btn = gr.Button("🚀 Start Training")
+        train_log = gr.Textbox(label="Training log", lines=10)
+        train_status = gr.Textbox(label="Status")
+
+        def train_click(bm, e, bs, lrn, wu, wdec, mi, mo, ic, oc, er, ga, fp):
+            log, status = run_training(bm, e, bs, lrn, wu, wdec, mi, mo, ic, oc, er, ga, fp)
+            return log, status
+
+        train_btn.click(train_click, [base_model, epochs, batch_size, lr, warmup, wd, max_in, max_out, in_col, out_col, eval_ratio, grad_accum, use_fp16], [train_log, train_status])
+
+    with gr.Tab("Demo"):
+        gr.Markdown("After training, this tab uses the local **checkpoint**.")
+        inp = gr.Textbox(label="Input Text", lines=10, placeholder="Paste text here...")
+        max_new_tokens = gr.Slider(16, 256, value=128, step=8, label="Max new tokens")
+        temperature = gr.Slider(0, 1.0, value=0.0, step=0.1, label="Temperature")
+        topp = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
+        btn = gr.Button("Summarize")
+        out = gr.Textbox(label="Summary", lines=10)
+
+        pipe_holder = {"pipe": None}
+
+        def summarize_click(text, max_new_tokens, temperature, top_p):
+            if pipe_holder["pipe"] is None:
+                pipe_holder["pipe"] = make_pipe_from_checkpoint()
+            gen = pipe_holder["pipe"](
+                f"summarize: {text}",
+                max_new_tokens=int(max_new_tokens),
+                do_sample=float(temperature)>0,
+                temperature=float(temperature),
+                top_p=float(top_p)
+            )
+            return gen[0]["generated_text"]
+
+        btn.click(summarize_click, [inp, max_new_tokens, temperature, topp], [out])
+
+if __name__ == "__main__":
+    demo.launch()