upload

README.md

@@ -0,0 +1,69 @@
+---
+language: en
+datasets:
+- sentence-transformers/reddit-title-body
+- sentence-transformers/embedding-training-data
+widget:
+- text: "Python is an interpreted, high-level and general-purpose programming language. Python's design philosophy emphasizes code readability with its notable use of significant whitespace. Its language constructs and object-oriented approach aim to help programmers write clear, logical code for small and large-scale projects."
+
+license: apache-2.0
+---
+
+# doc2query/all-t5-base-v1
+
+This is a [doc2query](https://arxiv.org/abs/1904.08375) model based on T5 (also known as [docT5query](https://cs.uwaterloo.ca/~jimmylin/publications/Nogueira_Lin_2019_docTTTTTquery-v2.pdf)).
+
+It can be used for:
+- **Document expansion**: You generate for your paragraphs 20-40 queries and index the paragraphs and the generates queries in a standard BM25 index like Elasticsearch, OpenSearch, or Lucene. The generated queries help to close the lexical gap of lexical search, as the generate queries contain synonyms. Further, it re-weights words giving important words a higher weight even if they appear seldomn in a paragraph. In our [BEIR](https://arxiv.org/abs/2104.08663) paper we showed that BM25+docT5query is a powerful search engine. In the [BEIR repository](https://github.com/UKPLab/beir) we have an example how to use docT5query with Pyserini.
+- **Domain Specific Training Data Generation**: It can be used to generate training data to learn an embedding model. On [SBERT.net](https://www.sbert.net/examples/unsupervised_learning/query_generation/README.html) we have an example how to use the model to generate (query, text) pairs for a given collection of unlabeled texts. These pairs can then be used to train powerful dense embedding models.
+
+## Usage
+```python
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+model_name = 'doc2query/all-t5-base-v1'
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name)
+
+text = "Python is an interpreted, high-level and general-purpose programming language. Python's design philosophy emphasizes code readability with its notable use of significant whitespace. Its language constructs and object-oriented approach aim to help programmers write clear, logical code for small and large-scale projects."
+
+
+input_ids = tokenizer.encode(text, max_length=384, truncation=True, return_tensors='pt')
+outputs = model.generate(
+    input_ids=input_ids,
+    max_length=64,
+    do_sample=True,
+    top_p=0.95,
+    num_return_sequences=5)
+
+print("Text:")
+print(text)
+
+print("\nGenerated Queries:")
+for i in range(len(outputs)):
+    query = tokenizer.decode(outputs[i], skip_special_tokens=True)
+    print(f'{i + 1}: {query}')
+```
+
+**Note:** `model.generate()` is non-deterministic. It produces different queries each time you run it.
+
+## Training
+This model fine-tuned [google/t5-v1_1-base](https://huggingface.co/google/t5-v1_1-base) for 570k training steps. For the  training script, see the `train_script.py` in this repository.
+
+The input-text was truncated to 384 word pieces. Output text was generated up to 64 word pieces. 
+
+This model was trained on a large collection of datasets. For the exact datasets names and weights see the `data_config.json` in this repository. Most of the datasets are available at [https://huggingface.co/sentence-transformers](https://huggingface.co/sentence-transformers).
+
+The datasets include besides others:
+- (title, body) pairs from [Reddit](https://huggingface.co/datasets/sentence-transformers/reddit-title-body)
+- (title, body) pairs and (title, answer) pairs from StackExchange and Yahoo Answers!
+- (title, review) pairs from Amazon reviews
+- (query, paragraph) pairs from MS MARCO, NQ, and GooAQ 
+- (question, duplicate_question) from Quora and WikiAnswers
+- (title, abstract) pairs from S2ORC
+
+## Prefix
+
+This model was trained **without a prefix**. In contrast to [doc2query/all-with_prefix-t5-base-v1](https://huggingface.co/doc2query/all-with_prefix-t5-base-v1) you cannot specify what type of transformation (answer2question, review2title) etc. you will have. This can lead to a mixture of output values.
+
+

config.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "google/t5-v1_1-base",
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 768,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "gradient_checkpointing": false,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "num_decoder_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.10.2",
+  "use_cache": true,
+  "vocab_size": 32128
+}

data_config.json

@@ -0,0 +1,51 @@
+{
+    "text2reddit": {
+        "weight": 30,
+        "files": {"/home/reddit/reddit-title-body/reddit_title_text_2010.jsonl.gz": 1,
+             "/home/reddit/reddit-title-body/reddit_title_text_2011.jsonl.gz":  1,
+             "/home/reddit/reddit-title-body/reddit_title_text_2012.jsonl.gz":  3,
+             "/home/reddit/reddit-title-body/reddit_title_text_2013.jsonl.gz":  5,
+             "/home/reddit/reddit-title-body/reddit_title_text_2014.jsonl.gz":  8,
+             "/home/reddit/reddit-title-body/reddit_title_text_2015.jsonl.gz":  11,
+             "/home/reddit/reddit-title-body/reddit_title_text_2016.jsonl.gz":  12,
+             "/home/reddit/reddit-title-body/reddit_title_text_2017.jsonl.gz":  13,
+             "/home/reddit/reddit-title-body/reddit_title_text_2018.jsonl.gz":  15,
+             "/home/reddit/reddit-title-body/reddit_title_text_2019.jsonl.gz":  19,
+             "/home/reddit/reddit-title-body/reddit_title_text_2020.jsonl.gz":  23,
+             "/home/reddit/reddit-title-body/reddit_title_text_2021.jsonl.gz":  12
+             }
+        
+    },
+    "question2title": {
+        "weight": 30,
+        "files": {"/home/sbert_pretrained_models/datasets/embedding-training-data/yahoo_answers_title_question.jsonl.gz": 40,
+            "/home/stackexchange_archive/jsonl/skeptics.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/writers.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/astronomy.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/vi.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/cstheory.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/engineering.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/french.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/economics.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/anime.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/islam.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/expressionengine.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/politics.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/history.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/christianity.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/boardgames.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/jsonl/civicrm.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/craftcms.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/hinduism.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/networkengineering.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/german.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/philosophy.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/gardening.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/space.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/bicycles.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/quant.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/puzzling.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/law.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/arduino.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/aviation.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/softwarerecs.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/movies.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/music.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/emacs.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/dsp.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/japanese.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/mechanics.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/crypto.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/cooking.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/photo.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/workplace.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/jsonl/biology.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/bitcoin.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/worldbuilding.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/datascience.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/ux.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/webapps.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/graphicdesign.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/raspberrypi.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/money.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/judaism.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/ethereum.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/academia.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/chemistry.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/webmasters.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/meta.stackoverflow.com.jsonl.gz": 3, "/home/stackexchange_archive/jsonl/cs.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/jsonl/travel.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/jsonl/rpg.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/jsonl/codereview.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/jsonl/gamedev.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/jsonl/android.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/jsonl/softwareengineering.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/jsonl/security.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/jsonl/diy.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/jsonl/scifi.stackexchange.com.jsonl.gz": 6, "/home/stackexchange_archive/jsonl/mathematica.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/jsonl/drupal.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/jsonl/blender.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/jsonl/dba.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/jsonl/ell.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/jsonl/meta.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/jsonl/gaming.stackexchange.com.jsonl.gz": 8, "/home/stackexchange_archive/jsonl/sharepoint.stackexchange.com.jsonl.gz": 8, "/home/stackexchange_archive/jsonl/magento.stackexchange.com.jsonl.gz": 9, "/home/stackexchange_archive/jsonl/wordpress.stackexchange.com.jsonl.gz": 9, "/home/stackexchange_archive/jsonl/salesforce.stackexchange.com.jsonl.gz": 9, "/home/stackexchange_archive/jsonl/english.stackexchange.com.jsonl.gz": 10, "/home/stackexchange_archive/jsonl/apple.stackexchange.com.jsonl.gz": 10, "/home/stackexchange_archive/jsonl/mathoverflow.net.jsonl.gz": 10, "/home/stackexchange_archive/jsonl/gis.stackexchange.com.jsonl.gz": 11, "/home/stackexchange_archive/jsonl/electronics.stackexchange.com.jsonl.gz": 12, "/home/stackexchange_archive/jsonl/physics.stackexchange.com.jsonl.gz": 15, "/home/stackexchange_archive/jsonl/stats.stackexchange.com.jsonl.gz": 15, "/home/stackexchange_archive/jsonl/unix.stackexchange.com.jsonl.gz": 16, "/home/stackexchange_archive/jsonl/tex.stackexchange.com.jsonl.gz": 17, "/home/stackexchange_archive/jsonl/serverfault.com.jsonl.gz": 23, "/home/stackexchange_archive/jsonl/askubuntu.com.jsonl.gz": 29, "/home/stackexchange_archive/jsonl/superuser.com.jsonl.gz": 36, "/home/stackexchange_archive/jsonl/small_stackexchanges.jsonl.gz": 37, "/home/stackexchange_archive/jsonl/math.stackexchange.com.jsonl.gz": 83, "/home/stackexchange_archive/jsonl/stackoverflow.com-Posts.jsonl.gz": 83
+            }
+    },
+    "answer2question": {
+        "weight": 30,
+        "files": {"/home/sbert_pretrained_models/datasets/embedding-training-data/yahoo_answers_title_answer.jsonl.gz": 80, "/home/sbert_pretrained_models/datasets/embedding-training-data/amazon-qa.jsonl.gz": 80,
+            "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/islam.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/anime.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/french.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/civicrm.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/expressionengine.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/history.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/politics.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/craftcms.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/christianity.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/softwarerecs.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/boardgames.stackexchange.com.jsonl.gz": 1, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/networkengineering.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/space.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/quant.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/philosophy.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/gardening.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/german.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/bicycles.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/law.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/arduino.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/emacs.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/dsp.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/puzzling.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/movies.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/mechanics.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/aviation.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/biology.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/crypto.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/music.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/datascience.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/japanese.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/bitcoin.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/cooking.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/photo.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/workplace.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/meta.stackoverflow.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/raspberrypi.stackexchange.com.jsonl.gz": 2, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/webapps.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/judaism.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/ethereum.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/worldbuilding.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/chemistry.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/graphicdesign.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/ux.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/money.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/cs.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/webmasters.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/academia.stackexchange.com.jsonl.gz": 3, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/travel.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/android.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/gamedev.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/rpg.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/codereview.stackexchange.com.jsonl.gz": 4, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/softwareengineering.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/security.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/diy.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/blender.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/scifi.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/mathematica.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/meta.stackexchange.com.jsonl.gz": 5, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/drupal.stackexchange.com.jsonl.gz": 6, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/dba.stackexchange.com.jsonl.gz": 6, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/ell.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/magento.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/sharepoint.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/gaming.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/wordpress.stackexchange.com.jsonl.gz": 7, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/mathoverflow.net.jsonl.gz": 8, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/salesforce.stackexchange.com.jsonl.gz": 8, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/apple.stackexchange.com.jsonl.gz": 8, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/gis.stackexchange.com.jsonl.gz": 9, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/english.stackexchange.com.jsonl.gz": 9, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/stats.stackexchange.com.jsonl.gz": 10, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/electronics.stackexchange.com.jsonl.gz": 11, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/physics.stackexchange.com.jsonl.gz": 12, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/unix.stackexchange.com.jsonl.gz": 13, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/tex.stackexchange.com.jsonl.gz": 15, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/serverfault.com.jsonl.gz": 20, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/askubuntu.com.jsonl.gz": 22, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/superuser.com.jsonl.gz": 30, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/small_stackexchanges.jsonl.gz": 38, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/math.stackexchange.com.jsonl.gz": 83, "/home/stackexchange_archive/stackexchange_extracted/TitleAnswer/stackoverflow.com-Posts.jsonl.gz": 83} 
+    },
+    "abstract2title": {
+        "weight": 10,
+        "files": {"/home/sbert_pretrained_models/datasets/embedding-training-data/S2ORC_title_abstract.jsonl.gz": 100}
+    },
+    "review2title": {
+        "weight": 10,
+        "files": {"/home/sbert_pretrained_models/datasets/embedding-training-data/amazon_review_2018.jsonl.gz": 100}
+    },
+    "news2title": {
+        "weight": 5,
+        "files": {"/home/sbert_pretrained_models/datasets/embedding-training-data/agnews.jsonl.gz": 50, "/home/sbert_pretrained_models/datasets/embedding-training-data/ccnews_title_text.jsonl.gz": 50, "/home/sbert_pretrained_models/datasets/embedding-training-data/xsum.jsonl.gz": 25}
+    },
+    "text2query": {
+        "weight": 10,
+        "files":  {"/home/sbert_pretrained_models/datasets/embedding-training-data/msmarco-triplets.jsonl.gz": 33, "/home/sbert_pretrained_models/datasets/embedding-training-data/gooaq_pairs.jsonl.gz": 66, "/home/sbert_pretrained_models/datasets/embedding-training-data/NQ-train_pairs.jsonl.gz": 6} 
+    },
+    "question2question": {
+        "weight": 10,
+        "files": {"/home/sbert_pretrained_models/datasets/embedding-training-data/WikiAnswers.jsonl.gz": 95, "/home/sbert_pretrained_models/datasets/embedding-training-data/quora_duplicates.jsonl.gz": 5}
+    }
+
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f44600793a9a002ab1a0d8d956b8c00f99adc0b9587e77485b7724de7a968f61
+size 990445401

special_tokens_map.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>", "additional_special_tokens": ["<extra_id_0>", "<extra_id_1>", "<extra_id_2>", "<extra_id_3>", "<extra_id_4>", "<extra_id_5>", "<extra_id_6>", "<extra_id_7>", "<extra_id_8>", "<extra_id_9>", "<extra_id_10>", "<extra_id_11>", "<extra_id_12>", "<extra_id_13>", "<extra_id_14>", "<extra_id_15>", "<extra_id_16>", "<extra_id_17>", "<extra_id_18>", "<extra_id_19>", "<extra_id_20>", "<extra_id_21>", "<extra_id_22>", "<extra_id_23>", "<extra_id_24>", "<extra_id_25>", "<extra_id_26>", "<extra_id_27>", "<extra_id_28>", "<extra_id_29>", "<extra_id_30>", "<extra_id_31>", "<extra_id_32>", "<extra_id_33>", "<extra_id_34>", "<extra_id_35>", "<extra_id_36>", "<extra_id_37>", "<extra_id_38>", "<extra_id_39>", "<extra_id_40>", "<extra_id_41>", "<extra_id_42>", "<extra_id_43>", "<extra_id_44>", "<extra_id_45>", "<extra_id_46>", "<extra_id_47>", "<extra_id_48>", "<extra_id_49>", "<extra_id_50>", "<extra_id_51>", "<extra_id_52>", "<extra_id_53>", "<extra_id_54>", "<extra_id_55>", "<extra_id_56>", "<extra_id_57>", "<extra_id_58>", "<extra_id_59>", "<extra_id_60>", "<extra_id_61>", "<extra_id_62>", "<extra_id_63>", "<extra_id_64>", "<extra_id_65>", "<extra_id_66>", "<extra_id_67>", "<extra_id_68>", "<extra_id_69>", "<extra_id_70>", "<extra_id_71>", "<extra_id_72>", "<extra_id_73>", "<extra_id_74>", "<extra_id_75>", "<extra_id_76>", "<extra_id_77>", "<extra_id_78>", "<extra_id_79>", "<extra_id_80>", "<extra_id_81>", "<extra_id_82>", "<extra_id_83>", "<extra_id_84>", "<extra_id_85>", "<extra_id_86>", "<extra_id_87>", "<extra_id_88>", "<extra_id_89>", "<extra_id_90>", "<extra_id_91>", "<extra_id_92>", "<extra_id_93>", "<extra_id_94>", "<extra_id_95>", "<extra_id_96>", "<extra_id_97>", "<extra_id_98>", "<extra_id_99>"]}

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

tokenizer_config.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>", "extra_ids": 100, "additional_special_tokens": ["<extra_id_0>", "<extra_id_1>", "<extra_id_2>", "<extra_id_3>", "<extra_id_4>", "<extra_id_5>", "<extra_id_6>", "<extra_id_7>", "<extra_id_8>", "<extra_id_9>", "<extra_id_10>", "<extra_id_11>", "<extra_id_12>", "<extra_id_13>", "<extra_id_14>", "<extra_id_15>", "<extra_id_16>", "<extra_id_17>", "<extra_id_18>", "<extra_id_19>", "<extra_id_20>", "<extra_id_21>", "<extra_id_22>", "<extra_id_23>", "<extra_id_24>", "<extra_id_25>", "<extra_id_26>", "<extra_id_27>", "<extra_id_28>", "<extra_id_29>", "<extra_id_30>", "<extra_id_31>", "<extra_id_32>", "<extra_id_33>", "<extra_id_34>", "<extra_id_35>", "<extra_id_36>", "<extra_id_37>", "<extra_id_38>", "<extra_id_39>", "<extra_id_40>", "<extra_id_41>", "<extra_id_42>", "<extra_id_43>", "<extra_id_44>", "<extra_id_45>", "<extra_id_46>", "<extra_id_47>", "<extra_id_48>", "<extra_id_49>", "<extra_id_50>", "<extra_id_51>", "<extra_id_52>", "<extra_id_53>", "<extra_id_54>", "<extra_id_55>", "<extra_id_56>", "<extra_id_57>", "<extra_id_58>", "<extra_id_59>", "<extra_id_60>", "<extra_id_61>", "<extra_id_62>", "<extra_id_63>", "<extra_id_64>", "<extra_id_65>", "<extra_id_66>", "<extra_id_67>", "<extra_id_68>", "<extra_id_69>", "<extra_id_70>", "<extra_id_71>", "<extra_id_72>", "<extra_id_73>", "<extra_id_74>", "<extra_id_75>", "<extra_id_76>", "<extra_id_77>", "<extra_id_78>", "<extra_id_79>", "<extra_id_80>", "<extra_id_81>", "<extra_id_82>", "<extra_id_83>", "<extra_id_84>", "<extra_id_85>", "<extra_id_86>", "<extra_id_87>", "<extra_id_88>", "<extra_id_89>", "<extra_id_90>", "<extra_id_91>", "<extra_id_92>", "<extra_id_93>", "<extra_id_94>", "<extra_id_95>", "<extra_id_96>", "<extra_id_97>", "<extra_id_98>", "<extra_id_99>"], "model_max_length": 512, "name_or_path": "google/t5-v1_1-base", "special_tokens_map_file": "/root/.cache/huggingface/transformers/76bf19bfedb85afbe644966ca9ab7b0404d753a41bf601115bced39f825ffa9c.c94798918c92ded6aeef2d2f0e666d2cc4145eca1aa6e1336fde07f2e13e2f46", "sp_model_kwargs": {}, "tokenizer_class": "T5Tokenizer"}

train_script.py

@@ -0,0 +1,253 @@
+import argparse
+import logging
+from torch.utils.data import Dataset, IterableDataset
+import gzip
+import json
+from transformers import Seq2SeqTrainer, AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainingArguments
+import sys
+from datetime import datetime
+import torch
+import random
+from shutil import copyfile
+import os
+import wandb
+import re
+
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_name", default="google/t5-v1_1-base")
+parser.add_argument("--train_file", required=True)
+parser.add_argument("--epochs", default=1, type=int)
+parser.add_argument("--batch_size", default=16, type=int)
+parser.add_argument("--max_source_length", default=384, type=int)
+parser.add_argument("--max_target_length", default=64, type=int)
+parser.add_argument("--name", required=True)
+parser.add_argument("--train_size", default=100*1000*1000, type=int)
+parser.add_argument("--eval_size", default=10000, type=int)
+parser.add_argument("--fp16", default=False, action='store_true')
+parser.add_argument("--no_prefix", default=False, action='store_true')
+args = parser.parse_args()
+
+wandb.init(project="doc2query", name=f"{args.name}-{args.model_name}")
+
+
+class PairDataset:
+    def __init__(self, filepath):
+        self.filepath = filepath
+        self.examples = []
+
+    def __iter__(self):
+        with gzip.open(self.filepath, 'rt') as fIn:
+            for line in fIn:
+                example = self.get_example(json.loads(line))
+               
+                if example is not None:
+                    self.examples.append(example)
+                    yield example
+
+        while True:
+            random.shuffle(self.examples)
+            for ex in self.examples:
+                yield ex
+
+
+    def get_example(self, raw_example):
+        if isinstance(raw_example, dict): 
+            if 'set' in raw_example:
+                example = random.sample(raw_example['set'], 2)
+            elif 'query' in raw_example:
+                example = [raw_example['query'], random.choice(raw_example['pos'])]
+            else:
+                raise ValueError("Unknown format: "+str(raw_example))
+        else:
+            example = [raw_example[0], raw_example[1]]
+
+        return example
+
+
+
+
+class RedditTitleDataset(PairDataset):
+    def get_example(self, raw_example):
+        return [self.clean_title(raw_example['title']), raw_example['body']]
+
+
+    def clean_title(self, text):
+        text = text.replace("&", "&").strip()
+        if text.startswith("["):
+            text = re.sub("^\[[a-zA-Z0-9]+\]", "", text).strip()
+
+        if text.endswith("]"):
+            text = re.sub("\[[a-zA-Z0-9\.]+\]$", "", text).strip()
+
+        if text.startswith("/r"):
+            text = re.sub("^/[a-zA-Z0-9/]+[;,: \-]+", "", text).strip()
+
+        return text
+
+
+class StackExchangeTitleBodyDataset(PairDataset):
+    def get_example(self, raw_example):
+        return raw_example['texts']
+
+
+class MultiDataset(IterableDataset):
+    def __init__(self, train_config_path, num_samples):
+        self.num_samples = num_samples
+
+        with open(train_config_path) as fIn:
+            train_config = json.load(fIn)
+
+        self.categories = []
+        self.files = {}
+        self.file2dataset = {}
+        self.file2datasetIter = {}
+
+        for prefix in train_config:
+            self.categories.extend([prefix]*train_config[prefix]['weight'])
+            self.files[prefix] = []
+
+            for filename, weight in train_config[prefix]['files'].items():
+                self.files[prefix].extend([filename]*weight)
+                dataset = self.OpenDataset(filename)
+                self.file2dataset[filename] = dataset
+                self.file2datasetIter[filename] = iter(dataset)
+
+            random.shuffle(self.files[prefix])
+
+        random.shuffle(self.categories)
+
+
+
+        
+    def OpenDataset(self, filepath):
+        if 'reddit_title_text' in filepath:
+            dataset = RedditTitleDataset(filepath)
+        elif 'stackexchange_archive/jsonl' in filepath:
+            dataset = StackExchangeTitleBodyDataset(filepath)
+        else:
+            dataset = PairDataset(filepath)
+        return dataset
+            
+
+    def __len__(self):
+        return self.num_samples
+
+    def __iter__(self):
+        while True:
+            category = random.choice(self.categories)
+            filepath = random.choice(self.files[category])
+            dataset = self.file2datasetIter[filepath]
+            pair = next(dataset)
+
+            #Add prefix to the input
+            if not args.no_prefix:
+                pair[1] = category+": "+pair[1].strip()
+            yield pair
+
+    def delete_examples_cache(self):
+        for dataset in self.file2dataset.values():
+            dataset.examples = []
+
+
+
+def main():
+    ############ Model
+    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+    save_steps = 5000
+
+    output_dir = 'output/'+args.name+'-'+args.model_name.replace("/", "-")+'-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    print("Output dir:", output_dir)
+
+    # Write self to path
+    os.makedirs(output_dir, exist_ok=True)
+
+    copyfile(args.train_file, os.path.join(output_dir, 'data_config.json'))
+    train_script_path = os.path.join(output_dir, 'train_script.py')
+    copyfile(__file__, train_script_path)
+    with open(train_script_path, 'a') as fOut:
+        fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))
+
+    ####
+
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=output_dir,
+        fp16=args.fp16,
+        fp16_backend="amp",
+        per_device_train_batch_size=args.batch_size,
+        evaluation_strategy="steps",
+        save_steps=save_steps,
+        logging_steps=100,
+        eval_steps=save_steps, #logging_steps,
+        warmup_steps=1000,
+        save_total_limit=1,
+        num_train_epochs=args.epochs,
+        report_to="wandb",
+    )
+
+    ############ Arguments
+
+    ############ Load datasets
+
+
+    train_dataset = MultiDataset(args.train_file, args.train_size)
+    train_dataset_iter = iter(train_dataset)
+    eval_dataset = [next(train_dataset_iter) for _ in range(args.eval_size)]
+    train_dataset.delete_examples_cache()  #Make sure dev data is no re-used for training
+
+    for i in range(50):
+        print("Target:", eval_dataset[i][0])
+        print("Input:", eval_dataset[i][1])
+        print("\n\n===================\n\n")
+
+    print("Train dataset len:", len(train_dataset))
+
+   
+    def data_collator(examples):
+        targets = [row[0] for row in examples]
+        inputs = [row[1] for row in examples]
+        label_pad_token_id = -100
+
+        model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8 if training_args.fp16 else None)
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=args.max_target_length, padding=True, truncation=True, pad_to_multiple_of=8 if training_args.fp16 else None)
+
+        # replace all tokenizer.pad_token_id in the labels by -100 to ignore padding in the loss.
+        labels["input_ids"] = [
+            [(l if l != tokenizer.pad_token_id else label_pad_token_id) for l in label] for label in labels["input_ids"]
+        ]
+
+
+        model_inputs["labels"] = torch.tensor(labels["input_ids"])
+        return model_inputs
+
+    ## Define the trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator
+    )
+
+    ### Save the model
+    train_result = trainer.train()
+    trainer.save_model()
+    
+    
+if __name__ == "__main__":
+    main()
+
+# Script was called via:
+#python train_hf_trainer_prefix.py --train_file train_config.json --name all-datasets-v1-no_prefix --no_prefix