Update README.md

README.md

@@ -13,7 +13,12 @@ The **miCSE** language model is trained for sentence similarity computation. Tra
 The model intended to be used for encoding sentences or short paragraphs. Given an input text, the model produces a vector embedding, which captures the semantics. The embedding can be used for numerous tasks, e.g., **retrieval**, **clustering** or **sentence similarity** comparison (see example below). Sentence representations correspond to the embedding of the _**[CLS]**_ token.
 
 
+# Training data
+
+The model was trained on a random collection of **English** sentences from Wikipedia: [Training data file](https://huggingface.co/datasets/princeton-nlp/datasets-for-simcse/resolve/main/wiki1m_for_simcse.txt)
+
 # Model Usage
+## Example 1) - Sentence Similarity
 
 ```python
 from transformers import AutoTokenizer, AutoModel
@@ -61,9 +66,80 @@ print(f"Distance: {cos_sim[0,1].detach().item()}")
 
 ```
 
-# Training data
+## Example 2) - Clustering
 
-The model was trained on a random collection of **English** sentences from Wikipedia: [Training data file](https://huggingface.co/datasets/princeton-nlp/datasets-for-simcse/resolve/main/wiki1m_for_simcse.txt)
+```python
+from transformers import AutoTokenizer, AutoModel
+import torch.nn as nn
+import torch
+import numpy as np
+import tqdm
+from datasets import load_dataset
+import umap
+import umap.plot as umap_plot
+
+# Determine available hardware
+if torch.backends.mps.is_available():
+    device = torch.device("mps")
+elif torch.cuda.is_available():
+    device = torch.device("gpu")
+else:
+    device = torch.device("cpu")
+    
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("/Users/d065243/miCSE")
+model = AutoModel.from_pretrained("/Users/d065243/miCSE")
+
+# Load Twitter data for sentiment clustering
+dataset = load_dataset("tweet_eval", "sentiment")
+
+
+# Compute embeddings of the tweets
+
+# set batch size and maxium tweet token length
+batch_size = 50
+max_length = 128
+
+iterations = int(np.floor(len(dataset['train'])/batch_size))*batch_size
+
+embedding_stack = []
+classes = []
+for i in tqdm.notebook.tqdm(range(0,iterations,batch_size)):
+    # create batch
+    batch = tokenizer.batch_encode_plus(
+                    dataset['train'][i:i+batch_size]['text'],
+                    return_tensors='pt',
+                    padding=True,
+                    max_length=max_length,
+                    truncation=True
+                ).to(device)
+    classes = classes + dataset['train'][i:i+batch_size]['label'] 
+
+    # model inference without gradient
+    with torch.no_grad():
+        outputs = model(**batch, output_hidden_states=True, return_dict=True)
+        
+        embeddings = outputs.last_hidden_state[:,0]
+        
+       
+        embedding_stack.append( embeddings.cpu().clone() )
+
+embeddings = torch.vstack(embedding_stack)
+
+
+# Cluster embeddings in 2D with UMAP
+umap_model = umap.UMAP(n_neighbors=250,
+                    n_components=2,
+                    min_dist=1.0e-9,
+                    low_memory=True,
+                    angular_rp_forest=True,
+                    metric='cosine')
+umap_model.fit(embeddings)
+
+# Plot result
+umap_plot.points(umap_model, labels = np.array(classes),theme='fire')
+
+```
 
 # Benchmark