Upload 2 files

client2.py

@@ -0,0 +1,175 @@
+import argparse
+import warnings
+from collections import OrderedDict
+from rouge import Rouge
+import torch
+from torch.optim import AdamW
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from torch.utils.data import DataLoader
+from datasets import load_dataset
+import flwr as fl
+
+
+from huggingface_hub import notebook_login
+
+notebook_login()
+
+
+
+def load_data(node_id):
+    """Load dataset (training and eval)"""
+    dataset = load_dataset("lighteval/legal_summarization", "BillSum")
+    full_train_dataset = dataset["train"]
+    eval_datasetx = dataset["test"]
+
+    tokenizer = AutoTokenizer.from_pretrained("t5-small")
+
+    # Split the full training dataset into two halves
+    train_dataset_size = len(full_train_dataset)
+    train_dataset_1 = full_train_dataset.select(range(0, train_dataset_size // 100))
+    train_dataset_2 = full_train_dataset.select(range(train_dataset_size // 2, train_dataset_size))
+
+    eval_dataset = eval_datasetx.select(range(0, 100))
+
+    # Choose one half as the training data
+    train_dataset = train_dataset_1
+
+    train_dataset = train_dataset.map(
+        lambda x: tokenizer.prepare_seq2seq_batch(x["article"], x["summary"]),
+        batched=True,
+    )
+
+    eval_dataset = eval_dataset.map(
+        lambda x: tokenizer.prepare_seq2seq_batch(x["article"], x["summary"]),
+        batched=True,
+    )
+
+    trainloader = DataLoader(train_dataset, batch_size=4, collate_fn=lambda data: collate_fn(data, tokenizer))
+    evalloader = DataLoader(eval_dataset, batch_size=4, collate_fn=lambda data: collate_fn(data, tokenizer))
+
+    return trainloader, evalloader, eval_dataset
+
+
+def collate_fn(data, tokenizer):
+    """Collate function to convert data into tensors"""
+    # Initialize lists to store tokenized articles and summaries
+    tokenized_articles = []
+    tokenized_summaries = []
+
+    # Iterate over each dictionary in the list
+    for item in data:
+        # Tokenize the article and summary
+        tokenized_item = tokenizer(item["article"], item["summary"], truncation=True, padding=True, return_tensors="pt")
+
+        # Append tokenized article to the list
+        tokenized_articles.append(tokenized_item["input_ids"])
+
+        # Check if "labels" key is present in the tokenized item
+        if "labels" in tokenized_item and "labels" in tokenized_item:
+            # If "labels" key is present, append tokenized summary to the list
+            tokenized_summaries.append(tokenized_item["labels"])
+        else:
+            # If "labels" key is not present, use "input_ids" as a placeholder for the summary
+            # You may need to adjust this logic based on the tokenizer's behavior
+            tokenized_summaries.append(tokenized_item["input_ids"])
+
+    # Convert lists to tensors
+    tokenized_articles = torch.stack(tokenized_articles).squeeze(dim=1)  # Remove singleton dimension
+    tokenized_summaries = torch.stack(tokenized_summaries).squeeze(dim=1)  # Remove singleton dimension
+
+    return {"input_ids": tokenized_articles, "labels": tokenized_summaries}
+
+
+def train(net, trainloader, epochs):
+    optimizer = AdamW(net.parameters(), lr=5e-5)
+    net.train()
+    total_batches = len(trainloader)
+    print("Training started...")
+    for i, batch in enumerate(trainloader, start=1):
+        inputs = {k: v.to(torch.device("cuda")) for k, v in batch.items()}  # Move all tensors to GPU
+        labels = inputs.pop("labels", None)  # Remove labels from inputs
+        outputs = net(**inputs, labels=labels) if labels is not None else net(**inputs)
+        loss = outputs.loss
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+        # Print progress within the single epoch
+        print(f"\rBatch {i}/{total_batches} - Loss: {loss.item():.4f}", end="", flush=True)
+    print("\nTraining finished.")
+
+    return net.state_dict()
+
+
+def calculate_rouge(net, eval_dataset, tokenizer):
+    rouge = Rouge()
+    references = [example["summary"] for example in eval_dataset]
+
+    generated_summaries = []
+    for example in eval_dataset:
+        input_ids = tokenizer(example["article"], truncation=True, padding=True, return_tensors="pt")["input_ids"]
+        outputs = net.generate(input_ids.to("cuda"))
+        generated_summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated_summaries.append(generated_summary)
+
+    scores = rouge.get_scores(generated_summaries, references)
+    rouge_1 = scores[0]["rouge-1"]["f"]
+    rouge_2 = scores[0]["rouge-2"]["f"]
+    rouge_l = scores[0]["rouge-l"]["f"]
+
+    return rouge_1, rouge_2, rouge_l
+
+def main(node_id):
+    net = AutoModelForSeq2SeqLM.from_pretrained("t5-small").to("cuda")
+
+    trainloader, _, eval_dataset = load_data(node_id)
+
+    # Flower client
+    class PlaceholderClient(fl.client.NumPyClient):
+        def get_parameters(self, config):
+            return [val.cpu().numpy() for _, val in net.state_dict().items()]
+
+        def set_parameters(self, parameters):
+            params_dict = zip(net.state_dict().keys(), parameters)
+            state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
+            net.load_state_dict(state_dict, strict=True)
+
+        def fit(self, parameters, config):
+            self.set_parameters(parameters)
+            print("Training Started...")
+            final_state_dict = train(net, trainloader, epochs=1)
+            print("Training Finished.")
+            return self.get_parameters(config={}), len(trainloader), {}
+
+        def evaluate(self, parameters, config):
+            self.set_parameters(parameters)
+            tokenizer = AutoTokenizer.from_pretrained("t5-small")
+            rouge_1, rouge_2, rouge_l = calculate_rouge(net, eval_dataset, tokenizer)
+            print(f"ROUGE-1 Score: {rouge_1:.4f}")
+            print(f"ROUGE-2 Score: {rouge_2:.4f}")
+            print(f"ROUGE-L Score: {rouge_l:.4f}")
+            # Replace 0.0 with a tuple or list of three elements
+            return 0.0, len(eval_dataset), {
+                "rouge-1": float(rouge_1),
+                "rouge-2": float(rouge_2),
+                "rouge-l": float(rouge_l),
+            }
+
+    # Start client
+    fl.client.start_client(
+        server_address="127.0.0.1:8089", client=PlaceholderClient().to_client()
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Flower")
+    parser.add_argument(
+        "--node-id",
+        choices=list(range(3)),
+        required=True,
+        type=int,
+        help="Partition of the dataset divided into 1,000 iid partitions created "
+             "artificially.",
+    )
+    node_id = parser.parse_args().node_id
+    main(node_id)

server2.py

@@ -0,0 +1,13 @@
+import flwr as fl
+
+
+strategy = fl.server.strategy.FedAvg(
+    fraction_fit=1.0, fraction_evaluate=1.0,
+)
+
+# Start server
+fl.server.start_server(
+    server_address="0.0.0.0:8089",
+    config=fl.server.ServerConfig(num_rounds=1),
+    strategy=strategy,
+)