Merge pull request #42 from borisdayma/chore-clean

chore: cleanup repo
Former-commit-id: 9977d1dc821ac8be7eef928e1aa6e2aaacd2c5f7

README.md

@@ -1,6 +1,6 @@
 ---
 title: Dalle Mini
-emoji: 🎨
+emoji: 🥑
 colorFrom: red
 colorTo: blue
 sdk: gradio
@@ -12,13 +12,17 @@ pinned: false
 
 _Generate images from a text prompt_
 
-TODO: add some cool example
+<img src="img/logo.png" width="200">
+
+Our logo was generated with DALL-E mini by typing "logo of an armchair in the shape of an avocado".
+
+You can also create your own pictures with the demo (TODO: add link).
 
 ## Create my own images with the demo → Coming soon
 
 ## How does it work?
 
-Refer to [our report](https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA?accessToken=2ua7j8ebc810fuxyv49wbipmq3fb2e78yq3rvs5dy4wew07wwm2csdo8zcuyr14e).
+Refer to [our report](https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA).
 
 ## Development
 

dev/seq2seq/run_seq2seq_flax.py

@@ -83,6 +83,7 @@ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
 
 # Model hyperparameters, for convenience
+# TODO: the model has now it's own definition file and should be imported
 OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos
 OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos
 BOS_TOKEN_ID = 16384
@@ -217,7 +218,7 @@ class DataTrainingArguments:
         default=None, metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
     )
     predict_with_generate: bool = field(
-        default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}
+        default=False, metadata={"help": "Whether to use generate to calculate generative metrics."}
     )
     num_beams: Optional[int] = field(
         default=None,
@@ -376,9 +377,6 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    logger.warning(f"WARNING: eval_steps has been manually hardcoded")  # TODO: remove it later, convenient for now
-    training_args.eval_steps = 400
-
     if (
         os.path.exists(training_args.output_dir)
         and os.listdir(training_args.output_dir)
@@ -425,11 +423,10 @@ def main():
     # (the dataset will be downloaded automatically from the datasets Hub).
     #
     data_files = {}
-    logger.warning(f"WARNING: Datasets path have been manually hardcoded")  # TODO: remove it later, convenient for now
     if data_args.train_file is not None:
-        data_files["train"] = ["/data/CC3M/training-encoded.tsv", "/data/CC12M/encoded-train.tsv", "/data/YFCC/metadata_encoded.tsv"]
+        data_files["train"] = data_args.train_file
     if data_args.validation_file is not None:
-        data_files["validation"] = ["/data/CC3M/validation-encoded.tsv"]
+        data_files["validation"] = data_args.validation_file
     if data_args.test_file is not None:
         data_files["test"] = data_args.test_file
     dataset = load_dataset("csv", data_files=data_files, cache_dir=model_args.cache_dir, delimiter="\t")
@@ -608,35 +605,6 @@ def main():
             desc="Running tokenizer on prediction dataset",
         )
 
-    # Metric
-    #metric = load_metric("rouge")
-
-    def postprocess_text(preds, labels):
-        preds = [pred.strip() for pred in preds]
-        labels = [label.strip() for label in labels]
-
-        # rougeLSum expects newline after each sentence
-        preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
-        labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]
-
-        return preds, labels
-
-    def compute_metrics(preds, labels):
-        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
-        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
-
-        # Some simple post-processing
-        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
-
-        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
-        # Extract a few results from ROUGE
-        result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
-
-        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
-        result["gen_len"] = np.mean(prediction_lens)
-        result = {k: round(v, 4) for k, v in result.items()}
-        return result
-
     # Initialize our training
     rng = jax.random.PRNGKey(training_args.seed)
     rng, dropout_rng = jax.random.split(rng)
@@ -822,15 +790,8 @@ def main():
             # log metrics
             wandb_log(eval_metrics, step=global_step, prefix='eval')
 
-            # compute ROUGE metrics
-            rouge_desc = ""
-        #    if data_args.predict_with_generate:
-        #        rouge_metrics = compute_metrics(eval_preds, eval_labels)
-        #        eval_metrics.update(rouge_metrics)
-        #        rouge_desc = " ".join([f"Eval {key}: {value} |" for key, value in rouge_metrics.items()])
-
             # Print metrics and update progress bar
-            desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {rouge_desc})"
+            desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']})"
             epochs.write(desc)
             epochs.desc = desc
 
@@ -955,15 +916,8 @@ def main():
         pred_metrics = get_metrics(pred_metrics)
         pred_metrics = jax.tree_map(jnp.mean, pred_metrics)
 
-        # compute ROUGE metrics
-        rouge_desc = ""
-        if data_args.predict_with_generate:
-            rouge_metrics = compute_metrics(pred_generations, pred_labels)
-            pred_metrics.update(rouge_metrics)
-            rouge_desc = " ".join([f"Predict {key}: {value} |" for key, value in rouge_metrics.items()])
-
         # Print metrics
-        desc = f"Predict Loss: {pred_metrics['loss']} | {rouge_desc})"
+        desc = f"Predict Loss: {pred_metrics['loss']})"
         logger.info(desc)