Update README.md

README.md

@@ -1,199 +1,115 @@
 ---
-library_name: transformers
-tags: []
+license: apache-2.0
+base_model:
+- answerdotai/ModernBERT-large
+datasets:
+- codelion/optillm-router-dataset
 ---
 
-# Model Card for Model ID
-
-<!-- Provide a quick summary of what the model is/does. -->
-
-
-
-## Model Details
-
-### Model Description
-
-<!-- Provide a longer summary of what this model is. -->
-
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
-
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
-
-### Model Sources [optional]
-
-<!-- Provide the basic links for the model. -->
-
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-
-## Uses
-
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-
-### Direct Use
-
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
-
-[More Information Needed]
-
-### Downstream Use [optional]
-
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
-
-[More Information Needed]
-
-### Out-of-Scope Use
-
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
-
-[More Information Needed]
-
-## Bias, Risks, and Limitations
-
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-
-[More Information Needed]
-
-### Recommendations
-
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
-
-## How to Get Started with the Model
-
-Use the code below to get started with the model.
-
-[More Information Needed]
-
-## Training Details
-
-### Training Data
-
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
-
-[More Information Needed]
-
-### Training Procedure
-
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-
-#### Preprocessing [optional]
-
-[More Information Needed]
-
-
-#### Training Hyperparameters
-
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
-
-#### Speeds, Sizes, Times [optional]
-
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-
-[More Information Needed]
-
-## Evaluation
-
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-[More Information Needed]
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-[More Information Needed]
-
-### Results
-
-[More Information Needed]
-
-#### Summary
-
-
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-[More Information Needed]
-
-## More Information [optional]
-
-[More Information Needed]
-
-## Model Card Authors [optional]
-
-[More Information Needed]
-
-## Model Card Contact
-
-[More Information Needed]
+# How to use?
+
+This model is used in [optillm](https://github.com/codelion/optillm) to route between the various approaches based on the prompt. 
+
+To use the model with optillm you can just prepend `router` to the model name. E.g. if we set `router-gpt-4o-mini` as the model, it will use the `gpt-4o-mini` as the base model.
+
+Otherwise, refer to the code in [router-plugin](https://github.com/codelion/optillm/blob/main/optillm/plugins/router_plugin.py) to see how to use this model for classification.
+
+# Usage
+
+To use the model directly you will need to use our `OptILMClassifier` class as we added additional layers to the base model. The additional
+`effort_encoder` is used to take into account the number of tokens a given approach consumes. Also, note
+the mapping of the returned index to the `APPROACHES` list as shown below. 
+
+```python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from safetensors.torch import load_model
+from transformers import AutoTokenizer, AutoModel
+
+# Constants
+MAX_LENGTH = 1024
+APPROACHES = ["none", "mcts", "bon", "moa", "rto", "z3", "self_consistency", "pvg", "rstar", "cot_reflection", "plansearch", "leap", "re2"]
+BASE_MODEL = "answerdotai/ModernBERT-large"
+OPTILLM_MODEL_NAME = "codelion/optillm-modernbert-large"
+
+class OptILMClassifier(nn.Module):
+    def __init__(self, base_model, num_labels):
+        super().__init__()
+        self.base_model = base_model
+        self.effort_encoder = nn.Sequential(
+            nn.Linear(1, 64),
+            nn.ReLU(),
+            nn.Linear(64, 64),
+            nn.ReLU()
+        )
+        self.classifier = nn.Linear(base_model.config.hidden_size + 64, num_labels)
+
+    def forward(self, input_ids, attention_mask, effort):
+        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.last_hidden_state[:, 0]  # Shape: (batch_size, hidden_size)
+        effort_encoded = self.effort_encoder(effort.unsqueeze(1))  # Shape: (batch_size, 64)
+        combined_input = torch.cat((pooled_output, effort_encoded), dim=1)
+        logits = self.classifier(combined_input)
+        return logits
+
+def load_optillm_model():
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+    # Load the base model
+    base_model = AutoModel.from_pretrained(BASE_MODEL)
+    # Create the OptILMClassifier
+    model = OptILMClassifier(base_model, num_labels=len(APPROACHES))  
+    model.to(device)
+    # Download the safetensors file
+    safetensors_path = hf_hub_download(repo_id=OPTILLM_MODEL_NAME, filename="model.safetensors")
+    # Load the state dict from the safetensors file
+    load_model(model, safetensors_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(OPTILLM_MODEL_NAME)
+    return model, tokenizer, device
+
+def preprocess_input(tokenizer, system_prompt, initial_query):
+    combined_input = f"{system_prompt}\n\nUser: {initial_query}"
+    encoding = tokenizer.encode_plus(
+        combined_input,
+        add_special_tokens=True,
+        max_length=MAX_LENGTH,
+        padding='max_length',
+        truncation=True,
+        return_attention_mask=True,
+        return_tensors='pt'
+    )
+    return encoding['input_ids'], encoding['attention_mask']
+
+def predict_approach(model, input_ids, attention_mask, device, effort=0.7):
+    model.eval()
+    with torch.no_grad():
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+        effort_tensor = torch.tensor([effort], dtype=torch.float).to(device)
+        
+        logits = model(input_ids, attention_mask=attention_mask, effort=effort_tensor)
+        probabilities = F.softmax(logits, dim=1)
+        predicted_approach_index = torch.argmax(probabilities, dim=1).item()
+        confidence = probabilities[0][predicted_approach_index].item()
+    
+    return APPROACHES[predicted_approach_index], confidence
+```
+
+You can now use the `predict_approach` method to get the predicted approach as follows:
+
+```python
+# Load the trained model
+router_model, tokenizer, device = load_optillm_model()
+
+# Preprocess the input
+input_ids, attention_mask = preprocess_input(tokenizer, system_prompt, initial_query)
+
+# Predict the best approach
+predicted_approach, _ = predict_approach(router_model, input_ids, attention_mask, device)
+
+print(f"Router predicted approach: {predicted_approach}")
+```