Create app.py

app.py

@@ -0,0 +1,120 @@
+import gradio as gr
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForSeq2Seq,
+)
+from datasets import load_dataset, Dataset
+import random
+
+# -----------------------------
+# Load Base Model
+# -----------------------------
+model_name = "PerceptronAI/Isaac-0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+
+# -----------------------------
+# Load Datasets
+# -----------------------------
+print("📥 Loading datasets...")
+
+pii_ds = load_dataset("ai4privacy/pii-masking-300k")
+cnn_ds = load_dataset("abisee/cnn_dailymail", "1.0.0")
+
+try:
+    docqa_ds = load_dataset("vidore/syntheticDocQA_energy_train")
+except Exception as e:
+    print("⚠️ Skipping docQA dataset (requires login):", e)
+    docqa_ds = None
+
+# -----------------------------
+# Build Training Samples
+# -----------------------------
+def make_pairs_pii(example):
+    return {"input": example["text"], "output": example["masked_text"]}
+
+def make_pairs_cnn(example):
+    return {"input": example["article"], "output": example["highlights"]}
+
+pii_pairs = pii_ds["train"].map(make_pairs_pii).select(range(1000))   # small subset
+cnn_pairs = cnn_ds["train"].map(make_pairs_cnn).select(range(1000))
+
+pairs = []
+pairs.extend(pii_pairs)
+pairs.extend(cnn_pairs)
+
+if docqa_ds is not None:
+    def make_pairs_docqa(example):
+        return {"input": example["question"], "output": example["answer"]}
+    docqa_pairs = docqa_ds["train"].map(make_pairs_docqa).select(range(1000))
+    pairs.extend(docqa_pairs)
+
+dataset = Dataset.from_list(pairs)
+
+# -----------------------------
+# Tokenization
+# -----------------------------
+def tokenize(batch):
+    inputs = tokenizer(batch["input"], truncation=True, padding="max_length", max_length=256)
+    outputs = tokenizer(batch["output"], truncation=True, padding="max_length", max_length=256)
+    inputs["labels"] = outputs["input_ids"]
+    return inputs
+
+tokenized_dataset = dataset.map(tokenize, batched=True)
+
+# -----------------------------
+# Training
+# -----------------------------
+training_args = TrainingArguments(
+    output_dir="./cass2.0",
+    overwrite_output_dir=True,
+    num_train_epochs=1,
+    per_device_train_batch_size=2,
+    save_steps=100,
+    save_total_limit=2,
+    logging_steps=20,
+    learning_rate=5e-5,
+    fp16=True,
+)
+
+data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+print("🚀 Training Cass2.0...")
+trainer.train()
+print("✅ Training complete!")
+
+# -----------------------------
+# Simple Chat UI
+# -----------------------------
+from transformers import pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+def chat(message, history):
+    prompt = "".join([f"User: {m[0]}\nCass2.0: {m[1]}\n" for m in history])
+    prompt += f"User: {message}\nCass2.0:"
+    output = pipe(prompt, max_length=256, do_sample=True, temperature=0.7)[0]["generated_text"]
+    reply = output.split("Cass2.0:")[-1].strip()
+    history.append((message, reply))
+    return history, history
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Cass2.0 — Trained AI Assistant")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Type your message")
+    clear = gr.Button("Clear")
+
+    msg.submit(chat, [msg, chatbot], [chatbot, chatbot])
+    clear.click(lambda: None, None, chatbot)
+
+demo.launch()