Align sample code with github readme

README.md

@@ -121,45 +121,30 @@ Use the code below to get started with the model.
 
 ```python
 from utils.regression_head import RegressionHead
-from transformers import AutoModel, AutoTokenizer
-from torch import bfloat16, no_grad, nn
-from huggingface_hub import hf_hub_download
-
-
-embedding_model = AutoModel.from_pretrained(
-    'Snowflake/snowflake-arctic-embed-m-v2.0', 
-    trust_remote_code=True,
-    torch_dtype=bfloat16,
-    unpad_inputs=True,
-    add_pooling_layer=False,
-).cuda()
-
-file_path = hf_hub_download(
-    repo_id="Jackal-AI/JQL-Edu-Heads",
-    filename="checkpoints/edu-mistral-snowflake-balanced.ckpt",
-    repo_type="model"
-)
-
-regression_head = RegressionHead.load_from_checkpoint(
-    file_path, 
-    map_location='cuda'
-).to(bfloat16)
-
-
-tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m-v2.0')
-tokens = tokenizer(
-    text=["Rome was founded by Romulus and Remus in 753 BC.", "Huggingface is cool!"], 
-    padding='longest', 
-    truncation=True, 
-    return_tensors='pt'
-).to('cuda')
-
-with no_grad():
-    cls_token = embedding_model(**tokens).last_hidden_state[:, 0]
-    normalized_cls_token = nn.functional.normalize(cls_token, p=2, dim=1)
-    predicted_edu_score = regression_head(normalized_cls_token).squeeze(-1)
-    
-print('predicted edu scores:', predicted_edu_score)
+from transformers.utils.hub import cached_file
+from utils.embedder import get_embedder_instance
+import torch
+
+# load embedder
+device = 'cuda'
+embedder = get_embedder_instance('Snowflake/snowflake-arctic-embed-m-v2.0', device, torch.bfloat16)
+# load JQL Edu annotation heads
+regression_head_checkpoints = {
+                'Edu-JQL-Gemma-SF': cached_file('Jackal-AI/JQL-Edu-Heads', 'checkpoints/edu-gemma-snowflake-balanced.ckpt'),
+                'Edu-JQL-Mistral-SF': cached_file('Jackal-AI/JQL-Edu-Heads', 'checkpoints/edu-mistral-snowflake-balanced.ckpt'),
+                'Edu-JQL-Llama-SF': cached_file('Jackal-AI/JQL-Edu-Heads', 'checkpoints/edu-llama-snowflake-balanced.ckpt'),
+            }
+regression_heads = {}
+for name, path in regression_head_checkpoints.items():
+    regression_heads[name] = RegressionHead.load_from_checkpoint(path, map_location=device).to(torch.bfloat16)
+
+# Given a single document
+doc = 'Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua'
+embeddings = embedder.embed([doc])
+scores = {}
+with torch.no_grad():
+    for name, regression_head in regression_heads.items():
+        scores[f'score_{name}'] = regression_head(embeddings).cpu().squeeze(1)
 ```
 <!-- 
 ## Training Details