Update app.py

app.py

@@ -1,29 +1,15 @@
-import gradio as gr
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    Trainer,
-    TrainingArguments,
-    DataCollatorForSeq2Seq,
-)
+import os
 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)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TrainingArguments, Trainer
+import evaluate
+import numpy as np
+import gradio as gr
 
 # -----------------------------
 # 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:
@@ -31,90 +17,107 @@ except Exception as e:
     docqa_ds = None
 
 # -----------------------------
-# Build Training Samples
+# Build Pairs from Datasets (Safe Version)
 # -----------------------------
-def make_pairs_pii(example):
-    return {"input": example["text"], "output": example["masked_text"]}
+pairs = []
 
-def make_pairs_cnn(example):
-    return {"input": example["article"], "output": example["highlights"]}
+def safe_map(dataset, input_keys, output_keys, name, limit=1000):
+    """
+    dataset: Hugging Face dataset split
+    input_keys: list of possible input column names
+    output_keys: list of possible output column names
+    name: dataset name (for logs)
+    limit: number of samples to select
+    """
+    available = dataset.column_names
+    chosen_in = next((k for k in input_keys if k in available), None)
+    chosen_out = next((k for k in output_keys if k in available), None)
 
-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))
+    if not chosen_in or not chosen_out:
+        print(f"⚠️ Skipping {name} (no matching columns). Available: {available}")
+        return []
 
-pairs = []
-pairs.extend(pii_pairs)
-pairs.extend(cnn_pairs)
+    print(f"✅ Using {name}: input='{chosen_in}', output='{chosen_out}'")
+
+    def make_pairs(example):
+        return {"input": example[chosen_in], "output": example[chosen_out]}
+
+    return dataset.map(make_pairs).select(range(min(limit, len(dataset))))
+
+pii_pairs = safe_map(pii_ds["train"], ["original", "text"], ["masked", "masked_text"], "PII")
+cnn_pairs = safe_map(cnn_ds["train"], ["article"], ["highlights", "summary"], "CNN/DailyMail")
 
 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)
+    docqa_pairs = safe_map(docqa_ds["train"], ["question"], ["answer"], "DocQA")
+else:
+    docqa_pairs = []
+
+pairs.extend(pii_pairs)
+pairs.extend(cnn_pairs)
+pairs.extend(docqa_pairs)
 
 dataset = Dataset.from_list(pairs)
 
 # -----------------------------
-# Tokenization
+# Model + Tokenizer
 # -----------------------------
-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
+model_name = "google/flan-t5-small"  # small, fast model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-tokenized_dataset = dataset.map(tokenize, batched=True)
+def tokenize_function(example):
+    model_inputs = tokenizer(example["input"], max_length=512, truncation=True)
+    labels = tokenizer(example["output"], max_length=128, truncation=True)
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
+
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
 # -----------------------------
 # Training
 # -----------------------------
+metric = evaluate.load("rouge")
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+    return {k: round(v * 100, 4) for k, v in result.items()}
+
 training_args = TrainingArguments(
-    output_dir="./cass2.0",
-    overwrite_output_dir=True,
+    output_dir="./results",
+    eval_strategy="no",
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
     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,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    logging_steps=10,
+    save_strategy="no"
 )
 
-data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_dataset,
+    train_dataset=tokenized_datasets,
+    eval_dataset=None,
     tokenizer=tokenizer,
-    data_collator=data_collator,
+    compute_metrics=compute_metrics
 )
 
-print("🚀 Training Cass2.0...")
 trainer.train()
-print("✅ Training complete!")
 
 # -----------------------------
-# Simple Chat UI
+# Gradio App
 # -----------------------------
-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()
+def generate_response(input_text):
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
+    outputs = model.generate(**inputs, max_new_tokens=128)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+demo = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Cass 2.0 Model")
+
+if __name__ == "__main__":
+    demo.launch()