Update README and config files - README.md

README.md

@@ -9,7 +9,10 @@ tags:
 - financial
 - mlm
 - roberta
+- embedding
 pipeline_tag: fill-mask
+tags_extended:
+- feature-extraction
 widget:
  - text: "<|cls|> Under the Migratory<|mask|> Treaty Act, the <|sep|>"
  - example: "<|cls|> This<|mask|> Credit Agreement is hereby <|sep|>"
@@ -20,7 +23,7 @@ date: '2024-11-07T00:00:00.000Z'
 
 # kl3m-doc-pico-001
 
-`kl3m-doc-pico-001` is a domain-specific masked language model (MLM) based on the RoBERTa architecture, specifically designed for legal and financial document analysis. With approximately 40M parameters, it provides a compact yet effective model for specialized NLP tasks.
+`kl3m-doc-pico-001` is a domain-specific masked language model (MLM) based on the RoBERTa architecture, specifically designed for legal and financial document analysis. With approximately 40M parameters, it provides a compact yet effective model for specialized NLP tasks in both fill-mask prediction and feature extraction for document embeddings.
 
 ## Model Details
 
@@ -28,8 +31,13 @@ date: '2024-11-07T00:00:00.000Z'
 - **Size**: 40M parameters
 - **Hidden Size**: 256
 - **Layers**: 8
+- **Attention Heads**: 8
+- **Intermediate Size**: 1024
+- **Max Position Embeddings**: 512
 - **Max Sequence Length**: 509
 - **Tokenizer**: [alea-institute/kl3m-004-128k-cased](https://huggingface.co/alea-institute/kl3m-004-128k-cased)
+- **Vector Dimension**: 256 (hidden_size)
+- **Pooling Strategy**: CLS token or mean pooling
 
 ## Use Cases
 
@@ -40,107 +48,273 @@ This model is particularly useful for:
 - Understanding legal citations and references
 - Filling in missing terms in legal documents
 - Feature extraction for downstream legal analysis tasks
+- Document similarity and retrieval tasks
+- Semantic search across legal and financial corpora
 
-As demonstrated in the examples, the model has strong knowledge of laws, regulations, and other domain-specific terminology.
+## Performance
+
+The model demonstrates strong performance on domain-specific tasks compared to general-purpose models of similar size. It shows particular strength in masked token prediction for legal and regulatory terminology.
+
+### Examples of Fill-Mask Performance
+
+1. Environmental Law:
+   - "Under the Migratory Bird Treaty Act, the..." → 93.4% confidence
+
+2. Financial Contracts:
+   - "This Farm Credit Agreement is hereby..." → 30.0% confidence
+
+## Standard Test Examples
+
+Using our standardized test examples for comparing embedding models:
+
+### 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. APPLICATION (0.0384)
+   2. PROCEDURES (0.0164)
+   3. HEARING (0.0150)
+   4. REGULATIONS (0.0119)
+   5. DEFINITIONS (0.0117)
+
+   Note: The contract-specialized models show stronger performance on this task, with kl3m-doc-pico-contracts-001 predicting "WARRANTIES" with higher confidence and kl3m-doc-nano-001 with 0.927 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.3148)
+   2. Date (0.2616)
+   3. Time (0.0220)
+   4. Order (0.0213)
+   5. Dates (0.0198)
+
+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.5241)
+   2. the (0.1538)
+   3. Disabilities (0.0708)
+   4. America (0.0228)
+   5. Truth (0.0147)
+
+### 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.711 | 0.675 |
+| Court Complaint vs. Credit Agreement | 0.847 | 0.833 |
+| Consumer Terms vs. Credit Agreement | 0.828 | 0.709 |
+
+Note how different pooling strategies affect similarity measurements. Mean pooling tends to capture more comprehensive document-level similarity, particularly between legal documents like complaints and agreements (0.833), while CLS token embeddings in this model show generally higher similarity values, particularly between Court Complaint and Credit Agreement (0.847).
 
 ## Usage
 
-You can use this model for masked language modeling with the following code:
+### Masked Language Modeling
+
+You can use this model for masked language modeling with the simple pipeline approach:
 
 ```python
-from transformers import AutoModelForMaskedLM, AutoTokenizer
-import torch
+from transformers import pipeline
+
+# Load the fill-mask pipeline with the model
+fill_mask = pipeline('fill-mask', model="alea-institute/kl3m-doc-pico-001")
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("alea-institute/kl3m-doc-pico-001")
-model = AutoModelForMaskedLM.from_pretrained("alea-institute/kl3m-doc-pico-001")
+# Example: Contract clause heading
+# 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)
 
-# Example 1: Legal domain - Environmental law
-text = "<|cls|> Under the Migratory<|mask|> Treaty Act, the <|sep|>"
-inputs = tokenizer(text, return_tensors="pt")
-outputs = model(**inputs)
+# Display predictions
+print("Top predictions:")
+for result in results:
+    print(f"- {result['token_str']} (score: {result['score']:.4f})")
 
-# Get top predictions for masked token
-masked_index = torch.where(inputs.input_ids[0] == tokenizer.mask_token_id)[0].item()
-probs = outputs.logits[0, masked_index].softmax(dim=0)
-top_5 = torch.topk(probs, 5)
+# Output:
+# Top predictions:
+# - APPLICATION (score: 0.0384)
+# - PROCEDURES (score: 0.0164)
+# - HEARING (score: 0.0150)
+# - REGULATIONS (score: 0.0119)
+# - DEFINITIONS (score: 0.0117)
+```
+
+You can also try other examples:
+
+```python
+# Example: 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("Example 1 - Top 5 predictions:")
-for i, (score, idx) in enumerate(zip(top_5.values, top_5.indices)):
-    token = tokenizer.decode(idx).strip()
-    print(f"{i+1}. {token} ({score.item():.3f})")
+# Display predictions
+print("Top predictions:")
+for result in results2[:5]:
+    print(f"- {result['token_str']} (score: {result['score']:.4f})")
 
 # Output:
-# Example 1 - Top 5 predictions:
-# 1. Bird (0.934)
-# 2. Species (0.028)
-# 3. Birds (0.006)
-# 4. Game (0.002)
-# 5. Fish (0.002)
-
-# Example 2: Financial domain - Credit agreement
-text = "<|cls|> This<|mask|> Credit Agreement is hereby <|sep|>"
-inputs = tokenizer(text, return_tensors="pt")
-outputs = model(**inputs)
-
-# Get top predictions for masked token
-masked_index = torch.where(inputs.input_ids[0] == tokenizer.mask_token_id)[0].item()
-probs = outputs.logits[0, masked_index].softmax(dim=0)
-top_5 = torch.topk(probs, 5)
-
-print("\nExample 2 - Top 5 predictions:")
-for i, (score, idx) in enumerate(zip(top_5.values, top_5.indices)):
-    token = tokenizer.decode(idx).strip()
-    print(f"{i+1}. {token} ({score.item():.3f})")
+# Top predictions:
+# - date (score: 0.3148)
+# - Date (score: 0.2616)
+# - Time (score: 0.0220)
+# - Order (score: 0.0213)
+# - Dates (score: 0.0198)
+```
+
+### Feature Extraction for Embeddings
 
+```python
+from transformers import pipeline
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+# Load the feature-extraction pipeline
+extractor = pipeline('feature-extraction', model="alea-institute/kl3m-doc-pico-001", return_tensors=True)
+
+# Example legal documents
+texts = [
+    # 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|>",
+    
+    # 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|>",
+    
+    # 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|>"
+]
+
+# 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.711 0.847]
+#  [0.711 1.    0.828]
+#  [0.847 0.828 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.675 0.833]
+#  [0.675 1.    0.709]
+#  [0.833 0.709 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}")
 # Output:
-# Example 2 - Top 5 predictions:
-# 1. Farm (0.300)
-# 2. Grant (0.036)
-# 3. Taxpayer (0.032)
-# 4. Negotiated (0.027)
-# 5. Credit (0.026)
+# Pairwise similarities:
+# Court Complaint vs. Consumer Terms:
+#   - CLS token: 0.7112
+#   - Mean pooling: 0.6749
+# Court Complaint vs. Credit Agreement:
+#   - CLS token: 0.8474
+#   - Mean pooling: 0.8331
+# Consumer Terms vs. Credit Agreement:
+#   - CLS token: 0.8276
+#   - Mean pooling: 0.7087
 ```
 
 ## Training
 
 The model was trained on a diverse corpus of legal and financial documents, ensuring high-quality performance in these domains. It leverages the KL3M tokenizer which provides 9-17% more efficient tokenization for domain-specific content than cl100k_base or the LLaMA/Mistral tokenizers.
 
+Training included both masked language modeling (MLM) objectives and attention to dense document representation for retrieval and classification tasks.
+
+## Intended Usage
+
+This model is intended for both:
+
+1. **Masked Language Modeling**: Filling in missing words/terms in legal and financial documents
+2. **Document Embedding**: Generating fixed-length vector representations for document similarity and classification
+
 ## 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
+- CLS token: `<|cls|>` (ID: 5) - Used for the beginning of input text
 - MASK token: `<|mask|>` (ID: 6) - Used to mark tokens for prediction
-- SEP token: `<|sep|>` (ID: 4) - Should be added at the end 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
 
+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
 
 While compact compared to larger language models, this model has some limitations:
 
-- Limited world knowledge compared to larger or more broadly-trained models
+- Limited parameter count (40M) means it captures less nuance than larger language models
 - Primarily focused on English legal and financial texts
 - Best suited for domain-specific rather than general-purpose tasks
+- Maximum sequence length of 512 tokens may require chunking for lengthy documents
+- Requires domain expertise to interpret results effectively
 
 ## 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
 
 If you use this model in your research, please cite:
 
 ```bibtex
-@misc{bommarito2025kl3m,
+@misc{kl3m-doc-pico-001,
+  author = {ALEA Institute},
+  title = {kl3m-doc-pico-001: A Domain-Specific Language Model for Legal and Financial Text Analysis},
+  year = {2024},
+  publisher = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/alea-institute/kl3m-doc-pico-001}}
+}
+
+@article{bommarito2025kl3m,
   title={KL3M Tokenizers: A Family of Domain-Specific and Character-Level Tokenizers for Legal, Financial, and Preprocessing Applications},
-  author={Bommarito II, Michael J. and Katz, Daniel Martin and Bommarito, Jillian},
+  author={Bommarito, Michael J and Katz, Daniel Martin and Bommarito, Jillian},
+  journal={arXiv preprint arXiv:2503.17247},
+  year={2025}
+}
+
+@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={2503.17247},
+  eprint={2504.07854},
   archivePrefix={arXiv},
   primaryClass={cs.CL}
 }
@@ -156,5 +330,6 @@ The KL3M model family is maintained by the [ALEA Institute](https://aleainstitut
  
 - Email: [email protected]
 - Website: https://aleainstitute.ai
+- GitHub: https://github.com/alea-institute/kl3m-model-research
 
-![https://aleainstitute.ai](https://aleainstitute.ai/images/alea-logo-ascii-1x1.png)
+![logo](https://aleainstitute.ai/images/alea-logo-ascii-1x1.png)