Update README and config files - README.md

README.md

@@ -44,81 +44,202 @@ This model is particularly useful for:
 - Feature extraction for downstream legal analysis tasks
 - Clustering documents by type, purpose, or content
 
-Document similarities from most similar to least similar:
+The moderately sized architecture makes this model a good balance between the lighter 41M parameter kl3m-doc-pico-001 model and larger language models, offering improved representation quality while remaining efficient for production deployments.
 
-1. Consumer Terms vs. Credit Agreement: 0.8092
-2. Consumer Terms vs. Financial Memo: 0.7604
-3. Court Complaint vs. Financial Memo: 0.7320
-4. Credit Agreement vs. Financial Memo: 0.7303
-5. Court Complaint vs. Credit Agreement: 0.6867
-6. Court Complaint vs. Consumer Terms: 0.6824
+## Standard Test Examples
 
-Note how documents of similar types (e.g., both contracts) show higher similarity despite different subject matter.
+Using our standardized test examples for comparing embedding models:
 
-The moderately sized architecture makes this model a good balance between the lighter 41M parameter kl3m-doc-pico-001 model and larger language models, offering improved representation quality while remaining efficient for production deployments.
+### Fill-Mask Results
+
+1. **Contract Clause Heading**:  
+   `"<|cls|> 8. REPRESENTATIONS AND<|mask|>. Each party hereby represents and warrants to the other party as of the date hereof as follows: <|sep|>"`
+
+   Top 5 predictions:
+   1. WARRANTIES (0.9271)
+   2. Warranties (0.0191)
+   3. REPRESENTATIONS (0.0156)
+   4. warrants (0.0123)
+   5. COVENANTS (0.0118)
+
+   Note: This model shows extremely strong performance on contract-specific language, predicting "WARRANTIES" with over 92% confidence.
+
+2. **Defined Term Example**:  
+   `"<|cls|> \"Effective<|mask|>\" means the date on which all conditions precedent set forth in Article V are satisfied or waived by the Administrative Agent. <|sep|>"`
+
+   Top 5 predictions:
+   1. Date (0.9601)
+   2. Time (0.0211)
+   3. date (0.0099)
+   4. Dates (0.0034)
+   5. Day (0.0006)
+
+3. **Regulation Example**:  
+   `"<|cls|> All transactions shall comply with the requirements set forth in the Truth in<|mask|> Act and its implementing Regulation Z. <|sep|>"`
+
+   Top 5 predictions:
+   1. Lending (0.9075)
+   2. the (0.0114)
+   3. Reinvestment (0.0083)
+   4. Disclosure (0.0036)
+   5. Credit (0.0026)
+
+Despite being designed primarily for feature extraction rather than masked language modeling, this model demonstrates excellent performance in predicting domain-specific terminology with high confidence.
+
+### Document Similarity Results
+
+Using the standardized document examples for embeddings:
+
+| Document Pair | Cosine Similarity (CLS token) | Cosine Similarity (Mean pooling) |
+|---------------|-------------------------------|----------------------------------|
+| Court Complaint vs. Consumer Terms | 0.682 | 0.721 |
+| Court Complaint vs. Credit Agreement | 0.687 | 0.770 |
+| Consumer Terms vs. Credit Agreement | 0.809 | 0.711 |
+
+The nano-sized model captures document similarities effectively with both CLS token and mean pooling strategies. With CLS tokens, it shows higher similarity between Consumer Terms and Credit Agreements (0.809), while with mean pooling, the highest similarity is observed between Court Complaints and Credit Agreements (0.770). This suggests the pooling strategy can significantly impact similarity measurements in this model with its moderate hidden size (512).
 
 ## Usage
 
-The primary use case for this model is feature extraction for document embedding and downstream classification tasks. Here's how to use it:
+The primary use case for this model is feature extraction for document embedding and downstream classification tasks.
+
+### Feature Extraction using Pipeline API
+
+The recommended way to use this model is through the Hugging Face pipeline API:
 
 ```python
-from transformers import AutoModel, AutoTokenizer
-import torch
+from transformers import pipeline
 import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("alea-institute/kl3m-doc-nano-001")
-model = AutoModel.from_pretrained("alea-institute/kl3m-doc-nano-001")
+# Load the feature-extraction pipeline
+extractor = pipeline('feature-extraction', model="alea-institute/kl3m-doc-nano-001", return_tensors=True)
 
-# Prepare diverse example texts from legal and financial domains
+# Example legal documents
 texts = [
-    # Example 1: Court Complaint
+    # Court Complaint
     "<|cls|> IN THE UNITED STATES DISTRICT COURT FOR THE EASTERN DISTRICT OF PENNSYLVANIA\n\nJOHN DOE,\nPlaintiff,\n\nvs.\n\nACME CORPORATION,\nDefendant.\n\nCIVIL ACTION NO. 21-12345\n\nCOMPLAINT\n\nPlaintiff John Doe, by and through his undersigned counsel, hereby files this Complaint against Defendant Acme Corporation, and in support thereof, alleges as follows: <|sep|>",
     
-    # Example 2: Consumer Terms and Conditions
+    # Consumer Terms
     "<|cls|> TERMS AND CONDITIONS\n\nLast Updated: April 10, 2025\n\nThese Terms and Conditions (\"Terms\") govern your access to and use of the Service. By accessing or using the Service, you agree to be bound by these Terms. If you do not agree to these Terms, you may not access or use the Service. These Terms constitute a legally binding agreement between you and the Company. <|sep|>",
     
-    # Example 3: Credit Agreement
-    "<|cls|> CREDIT AGREEMENT\n\nDated as of April 10, 2025\n\nAmong\n\nACME BORROWER INC.,\nas the Borrower,\n\nBANK OF FINANCE,\nas Administrative Agent,\n\nand\n\nTHE LENDERS PARTY HERETO\n\nThis CREDIT AGREEMENT (\"Agreement\") is entered into as of April 10, 2025, among ACME BORROWER INC., a Delaware corporation (the \"Borrower\"), each lender from time to time party hereto (collectively, the \"Lenders\"), and BANK OF FINANCE, as Administrative Agent. <|sep|>",
-    
-    # Example 4: Financial Memorandum
-    "<|cls|> MEMORANDUM\n\nTO: Investment Committee\nFROM: John Smith, Chief Investment Officer\nDATE: April 10, 2025\nSUBJECT: Q2 2025 Investment Strategy\n\nThis memorandum outlines our proposed investment strategy for Q2 2025. Based on current market conditions and economic indicators, we recommend the following asset allocation adjustments to the portfolio. <|sep|>"
+    # Credit Agreement
+    "<|cls|> CREDIT AGREEMENT\n\nDated as of April 10, 2025\n\nAmong\n\nACME BORROWER INC.,\nas the Borrower,\n\nBANK OF FINANCE,\nas Administrative Agent,\n\nand\n\nTHE LENDERS PARTY HERETO\n\nThis CREDIT AGREEMENT (\"Agreement\") is entered into as of April 10, 2025, among ACME BORROWER INC., a Delaware corporation (the \"Borrower\"), each lender from time to time party hereto (collectively, the \"Lenders\"), and BANK OF FINANCE, as Administrative Agent. <|sep|>"
 ]
 
-# Get embeddings for text classification or similarity
-inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
-outputs = model(**inputs)
-
-# Extract CLS token embeddings (for classification tasks)
-cls_embeddings = outputs.last_hidden_state[:, 0, :]
-print(f"CLS embeddings shape: {cls_embeddings.shape}")  # Should be [4, 512]
-
-# Calculate similarity between documents
-def cosine_similarity(a, b):
-    return torch.dot(a, b) / (torch.norm(a) * torch.norm(b))
-
-# Calculate and print similarities (sorted from highest to lowest)
-print("\nCosine similarities (from highest to lowest):")
-similarities = [
-    ("Consumer Terms vs. Credit Agreement", cosine_similarity(cls_embeddings[1], cls_embeddings[2]).item()),
-    ("Consumer Terms vs. Financial Memo", cosine_similarity(cls_embeddings[1], cls_embeddings[3]).item()),
-    ("Court Complaint vs. Financial Memo", cosine_similarity(cls_embeddings[0], cls_embeddings[3]).item()),
-    ("Credit Agreement vs. Financial Memo", cosine_similarity(cls_embeddings[2], cls_embeddings[3]).item()),
-    ("Court Complaint vs. Credit Agreement", cosine_similarity(cls_embeddings[0], cls_embeddings[2]).item()),
-    ("Court Complaint vs. Consumer Terms", cosine_similarity(cls_embeddings[0], cls_embeddings[1]).item())
-]
-
-# Sort by similarity score (descending)
-similarities.sort(key=lambda x: x[1], reverse=True)
-
-# Print sorted similarities
-for desc, score in similarities:
-    print(f"{desc}: {score:.4f}")
+# Strategy 1: CLS token embeddings
+cls_embeddings = []
+for text in texts:
+    features = extractor(text)
+    # Get the CLS token (first token) embedding
+    features_array = features[0].numpy() if hasattr(features[0], 'numpy') else features[0]
+    cls_embedding = features_array[0]
+    cls_embeddings.append(cls_embedding)
+
+# Calculate similarity between documents using CLS tokens
+cls_similarity = cosine_similarity(np.vstack(cls_embeddings))
+print("\nDocument similarity (CLS token):")
+print(np.round(cls_similarity, 3))
+# Output:
+# [[1.    0.682 0.687]
+#  [0.682 1.    0.809]
+#  [0.687 0.809 1.   ]]
+
+# Strategy 2: Mean pooling
+mean_embeddings = []
+for text in texts:
+    features = extractor(text)
+    # Average over all tokens
+    features_array = features[0].numpy() if hasattr(features[0], 'numpy') else features[0]
+    mean_embedding = np.mean(features_array, axis=0)
+    mean_embeddings.append(mean_embedding)
+
+# Calculate similarity using mean pooling
+mean_similarity = cosine_similarity(np.vstack(mean_embeddings))
+print("\nDocument similarity (Mean pooling):")
+print(np.round(mean_similarity, 3))
+# Output:
+# [[1.    0.721 0.77 ]
+#  [0.721 1.    0.711]
+#  [0.77  0.711 1.   ]]
+
+# Print pairwise similarities
+doc_names = ["Court Complaint", "Consumer Terms", "Credit Agreement"]
+print("\nPairwise similarities:")
+for i in range(len(doc_names)):
+    for j in range(i+1, len(doc_names)):
+        print(f"{doc_names[i]} vs. {doc_names[j]}:")
+        print(f"  - CLS token: {cls_similarity[i, j]:.4f}")
+        print(f"  - Mean pooling: {mean_similarity[i, j]:.4f}")
 
 # For document retrieval, you can index these embeddings in a vector database
 # such as FAISS, Pinecone, or Weaviate for efficient similarity search
 ```
 
+### Masked Language Modeling
+
+While primarily designed for feature extraction, the model also supports masked language modeling:
+
+```python
+from transformers import pipeline
+
+# Load the fill-mask pipeline
+fill_mask = pipeline('fill-mask', model="alea-institute/kl3m-doc-nano-001")
+
+# Example text with mask token
+# Note the mask token placement - directly adjacent to "AND" without space
+text = "<|cls|> 8. REPRESENTATIONS AND<|mask|>. Each party hereby represents and warrants to the other party as of the date hereof as follows: <|sep|>"
+results = fill_mask(text)
+
+# Display predictions
+print("Top predictions:")
+for result in results:
+    print(f"- {result['token_str']} (score: {result['score']:.4f})")
+
+# Output:
+# Top predictions:
+# - WARRANTIES (score: 0.9271)
+# - Warranties (score: 0.0191)
+# - REPRESENTATIONS (score: 0.0156)
+# - warrants (score: 0.0123)
+# - COVENANTS (score: 0.0118)
+```
+
+Here are some additional examples:
+
+```python
+# Example 2: Defined term
+text2 = "<|cls|> \"Effective<|mask|>\" means the date on which all conditions precedent set forth in Article V are satisfied or waived by the Administrative Agent. <|sep|>"
+results2 = fill_mask(text2)
+
+print("\nDefined Term Example - Top 5 predictions:")
+for i, result in enumerate(results2[:5]):
+    print(f"{i+1}. {result['token_str']} ({result['score']:.4f})")
+
+# Output:
+# Defined Term Example - Top 5 predictions:
+# 1. Date (0.9601)
+# 2. Time (0.0211)
+# 3. date (0.0099)
+# 4. Dates (0.0034)
+# 5. Day (0.0006)
+
+# Example 3: Regulatory context
+text3 = "<|cls|> All transactions shall comply with the requirements set forth in the Truth in<|mask|> Act and its implementing Regulation Z. <|sep|>"
+results3 = fill_mask(text3)
+
+print("\nRegulatory Example - Top 5 predictions:")
+for i, result in enumerate(results3[:5]):
+    print(f"{i+1}. {result['token_str']} ({result['score']:.4f})")
+
+# Output:
+# Regulatory Example - Top 5 predictions:
+# 1. Lending (0.9075)
+# 2. the (0.0114)
+# 3. Reinvestment (0.0083)
+# 4. Disclosure (0.0036)
+# 5. Credit (0.0026)
+```
+
 ## Training
 
 The model was trained on a diverse corpus of legal and financial documents, ensuring high-quality performance in these domains. It incorporates the following key aspects:
@@ -148,8 +269,9 @@ import torch
 tokenizer = AutoTokenizer.from_pretrained("alea-institute/kl3m-doc-nano-001")
 mlm_model = AutoModelForMaskedLM.from_pretrained("alea-institute/kl3m-doc-nano-001")
 
-# Example 1: Contract with masked entity type
-text1 = "<|cls|> This Credit Agreement is made and entered into as of January 1, 2025, by and between Acme Corporation, a Delaware <|mask|>, and First National Bank, as lender. <|sep|>"
+# Example 1: Contract clause with mask token immediately after word (no space)
+# Note the placement - directly after "AND" without a space
+text1 = "<|cls|> 8. REPRESENTATIONS AND<|mask|>. Each party hereby represents and warrants to the other party as of the date hereof as follows: <|sep|>"
 inputs1 = tokenizer(text1, return_tensors="pt")
 outputs1 = mlm_model(**inputs1)
 
@@ -158,20 +280,20 @@ masked_index1 = torch.where(inputs1.input_ids[0] == tokenizer.mask_token_id)[0].
 probs1 = outputs1.logits[0, masked_index1].softmax(dim=0)
 top_5_1 = torch.topk(probs1, 5)
 
-print("Credit Agreement Example - Top 5 predictions:")
+print("Contract Clause Example - Top 5 predictions:")
 for i, (score, idx) in enumerate(zip(top_5_1.values, top_5_1.indices)):
     token = tokenizer.decode(idx).strip()
     print(f"{i+1}. {token} ({score.item():.4f})")
 
 # Output:
-# 1. Corp. (0.1706)
-# 2. C. (0.1694)
-# 3. Bank. (0.1556)
-# 4. Inc. (0.0596)
-# 5. U.S. (0.0491)
-
-# Example 2: Regulatory document with masked term
-text2 = "<|cls|> Pursuant to Section 13(a) of the <|mask|> Exchange Act of 1934, the undersigned hereby certifies that the quarterly report fully complies with the requirements. <|sep|>"
+# 1. WARRANTIES (0.9271)
+# 2. Warranties (0.0191)
+# 3. REPRESENTATIONS (0.0156)
+# 4. warrants (0.0123) 
+# 5. COVENANTS (0.0118)
+
+# Example 2: Defined term with mask token immediately after word (no space)
+text2 = "<|cls|> \"Effective<|mask|>\" means the date on which all conditions precedent set forth in Article V are satisfied or waived by the Administrative Agent. <|sep|>"
 inputs2 = tokenizer(text2, return_tensors="pt")
 outputs2 = mlm_model(**inputs2)
 
@@ -180,20 +302,20 @@ masked_index2 = torch.where(inputs2.input_ids[0] == tokenizer.mask_token_id)[0].
 probs2 = outputs2.logits[0, masked_index2].softmax(dim=0)
 top_5_2 = torch.topk(probs2, 5)
 
-print("\nRegulatory Example - Top 5 predictions:")
+print("\nDefined Term Example - Top 5 predictions:")
 for i, (score, idx) in enumerate(zip(top_5_2.values, top_5_2.indices)):
     token = tokenizer.decode(idx).strip()
     print(f"{i+1}. {token} ({score.item():.4f})")
 
 # Output:
-# 1. U.S. (0.2473)
-# 2. Trad (0.2095)
-# 3. Securities (0.1251)
-# 4.  (0.0478)
-# 5. ## (0.0183)
-
-# Example 3: Corporate governance with masked term
-text3 = "<|cls|> The Board of Directors shall consist of not less than three nor more than fifteen <|mask|>, with the exact number to be fixed by resolution of the Board. <|sep|>"
+# 1. Date (0.9601)
+# 2. Time (0.0211)
+# 3. date (0.0099)
+# 4. Dates (0.0034)
+# 5. Day (0.0006)
+
+# Example 3: Regulatory example with mask token immediately after word (no space)
+text3 = "<|cls|> All transactions shall comply with the requirements set forth in the Truth in<|mask|> Act and its implementing Regulation Z. <|sep|>"
 inputs3 = tokenizer(text3, return_tensors="pt")
 outputs3 = mlm_model(**inputs3)
 
@@ -202,40 +324,46 @@ masked_index3 = torch.where(inputs3.input_ids[0] == tokenizer.mask_token_id)[0].
 probs3 = outputs3.logits[0, masked_index3].softmax(dim=0)
 top_5_3 = torch.topk(probs3, 5)
 
-print("\nCorporate Governance Example - Top 5 predictions:")
+print("\nRegulatory Example - Top 5 predictions:")
 for i, (score, idx) in enumerate(zip(top_5_3.values, top_5_3.indices)):
     token = tokenizer.decode(idx).strip()
     print(f"{i+1}. {token} ({score.item():.4f})")
 
 # Output:
-# 1. (15) (0.6526)
-# 2. (10) (0.0732)
-# 3. (14) (0.0627)
-# 4. (5) (0.0296)
-# 5. (3) (0.0244)
+# 1. Lending (0.9075)
+# 2. the (0.0114)
+# 3. Reinvestment (0.0083)
+# 4. Disclosure (0.0036)
+# 5. Credit (0.0026)
 ```
 
-These examples show that the model can predict contextually appropriate tokens across different legal and financial domains:
+These examples demonstrate the model's exceptional performance with domain-specific terminology and its ability to predict appropriate terms with very high confidence (over 90% in some cases). Note the careful placement of the mask token - immediately after a word without a space when we want to predict a word that would normally be adjacent to the previous word.
+
+These examples show that the model can predict contextually appropriate tokens across different legal and financial domains with extremely high confidence:
 
-1. In a credit agreement context, it correctly suggests entity suffixes like "Corp." and "Inc."
-2. In a regulatory context, it suggests appropriate terms like "Securities" and "U.S." 
-3. In corporate governance documents, it accurately predicts numeric parenthetical references like "(15)" that are common in bylaws and similar documents
+1. In a contract clause context, it correctly predicts "WARRANTIES" with over 92% confidence
+2. In a defined term context, it predicts "Date" with 96% confidence
+3. In a regulatory context, it correctly identifies "Lending" for the Truth in Lending Act with 90% confidence
 
 For applications primarily focused on masked language modeling, you may still want to consider models specifically optimized for that task.
 
 ## Special Tokens
 
-This model uses custom special tokens which must be used explicitly:
+This model includes the following special tokens:
 
-- CLS token: `<|cls|>` (ID: 5) - Should be added at the beginning of input text
-- SEP token: `<|sep|>` (ID: 4) - Should be added at the end of input text
+- CLS token: `<|cls|>` (ID: 5) - Used for the beginning of input text
+- SEP token: `<|sep|>` (ID: 4) - Used for the end of input text
 - PAD token: `<|pad|>` (ID: 2) - Used for padding sequences to a uniform length
 - BOS token: `<|start|>` (ID: 0) - Beginning of sequence
 - EOS token: `<|end|>` (ID: 1) - End of sequence
 - UNK token: `<|unk|>` (ID: 3) - Unknown token
 - MASK token: `<|mask|>` (ID: 6) - Used for masked language modeling tasks
 
-For best results, you should **explicitly add** the CLS and SEP tokens to your input text, as shown in the examples above. The CLS token embedding is particularly important as it's commonly used for classification tasks.
+Important usage notes:
+
+When using the MASK token for predictions, be aware that this model uses a **space-prefixed BPE tokenizer**. The <|mask|> token should be placed IMMEDIATELY after the previous token with NO space, because most tokens in this tokenizer have an initial space encoded within them. For example: `"word<|mask|>"` rather than `"word <|mask|>"`.
+   
+This space-aware placement is crucial for getting accurate predictions, as demonstrated in our test examples.
 
 ## Limitations
 
@@ -250,7 +378,7 @@ While providing a good balance between size and performance, this model has some
 ## References
 
 - [KL3M Tokenizers: A Family of Domain-Specific and Character-Level Tokenizers for Legal, Financial, and Preprocessing Applications](https://arxiv.org/abs/2503.17247)
-- [The KL3M Data Project: Copyright-Clean Training Resources for Large Language Models]() (Forthcoming)
+- [The KL3M Data Project: Copyright-Clean Training Resources for Large Language Models](https://arxiv.org/abs/2504.07854)
 
 ## Citation
 
@@ -267,6 +395,17 @@ If you use this model in your research, please cite:
 }
 ```
 
+```bibtex
+@misc{bommarito2025kl3mdata,
+  title={The KL3M Data Project: Copyright-Clean Training Resources for Large Language Models},
+  author={Bommarito II, Michael J. and Bommarito, Jillian and Katz, Daniel Martin},
+  year={2025},
+  eprint={2504.07854},
+  archivePrefix={arXiv},
+  primaryClass={cs.CL}
+}
+```
+
 ## License
 
 This model is licensed under [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/).