Update README.md

README.md

@@ -6,194 +6,191 @@ tags: []
 # Model Card for Model ID
 
 <!-- Provide a quick summary of what the model is/does. -->
-
-
+Classification model finetuned for prompt-model routing based on code prompt difficulty.
 
 ## Model Details
 
+Base model: answerdotai/ModernBERT-base
+
 ### 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
+- **Developed by:** [Christian @ Prime Intellect]
+- **Finetuned from model:** [answerdotai/ModernBERT-base]
 
-<!-- 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.
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+```
+from transformers import pipeline
+import torch
+
+# Load the model
+classifier = pipeline(
+    "text-classification",
+    model="cdreetz/modern-bert-router",
+    device=0 if torch.cuda.is_available() else -1
+)
+
+# Test easy problem
+easy_problem = """
+Write a function that returns the sum of two numbers.
+
+Example:
+Input: add(2, 3)
+Output: 5
+"""
+
+# Test hard problem
+hard_problem = """
+Given a binary tree, find the maximum path sum. The path may start and end at any node in the tree.
+A path is defined as any sequence of nodes from some starting node to any node in the tree along the parent-child connections.
+The path must contain at least one node and does not need to go through the root.
+
+Example:
+Input: root = [1,2,3]
+Output: 6 (2 -> 1 -> 3)
+"""
+
+# Run predictions
+result_easy = classifier(easy_problem)[0]
+result_hard = classifier(hard_problem)[0]
+
+print("Easy Problem:")
+print(f"  Difficulty: {result_easy['label']}")
+print(f"  Confidence: {result_easy['score']:.2%}\n")
+
+print("Hard Problem:")
+print(f"  Difficulty: {result_hard['label']}")
+print(f"  Confidence: {result_hard['score']:.2%}")
+```
 
-[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
+```
+# /// script
+# dependencies = [
+#   "chatan",
+#   "transformers",
+#   "datasets",
+#   "torch",
+#   "accelerate",
+#   "scikit-learn",
+#   "triton",
+#   "huggingface_hub"
+# ]
+# ///
+
+import os
+import asyncio
+import chatan as ch
+from datasets import Dataset as hf_ds
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+import numpy as np
+from sklearn.metrics import accuracy_score, f1_score
+import torch
+import triton
+from huggingface_hub import login
+
+#torch._dynamo.config.suppress_errors = True
+
+login()
+
+async def create_dataset():
+    gen = ch.async_generator("openai", os.getenv("OPENAI_API_KEY"))
+
+    ds = ch.async_dataset({
+        "difficulty": ch.sample.choice(["easy", "hard"]),
+        "text": gen("write a coding problem of difficulty {difficulty}")
+    })
+
+    df = await ds.generate(n=1000, max_concurrent_rows=500)
+    df['labels'] = df['difficulty'].map({"easy": 0, "hard": 1})
+    dataset = hf_ds.from_pandas(df[['text', 'labels']])
+    return dataset.train_test_split(test_size=0.2, seed=42)
+
+
+def train(dataset):
+    tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "answerdotai/ModernBERT-base",
+        num_labels=2,
+        label2id={"easy": 0, "hard": 1},
+        id2label={0: "easy", 1: "hard"},
+        problem_type="single_label_classification"
+    )
+
+
+    def tokenize(examples):
+        return tokenizer(
+            examples['text'], 
+            padding='max_length', 
+            truncation=True, 
+            max_length=512
+        )
+
+    tokenized = dataset.map(tokenize, batched=True, remove_columns=['text'])
+    tokenized.set_format('torch')
+
+    training_args = TrainingArguments(
+        output_dir="./modernbert-router",
+        eval_strategy="epoch",
+        num_train_epochs=3,
+        per_device_train_batch_size=16,
+        learning_rate=2e-5,
+        save_strategy="epoch",
+        load_best_model_at_end=True
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized["train"],
+        eval_dataset=tokenized["test"],
+        processing_class=tokenizer,
+        compute_metrics=lambda p: {
+            "accuracy": accuracy_score(p.label_ids, np.argmax(p.predictions, axis=1)),
+            "f1": f1_score(p.label_ids, np.argmax(p.predictions, axis=1))
+        }
+    )
+
+    print("starting training")
+    trainer.train()
+
+    model.save_pretrained("./modernbert-router")
+    tokenizer.save_pretrained("./modernbert-router")
+
+    print("donezo")
+
+    #model = AutoModelForSequenceClassification.from_pretrained("./modernbert-router")
+    #tokenizer = AutoTokenizer.from_pretrained("./modernbert-router")
+
+    #model.push_to_hub("cdreetz/modern-bert-router")
+    #tokenizer.push_to_hub("cdreetz/modern-bert-router")
+
+
+
+if __name__ == "__main__":
+    dataset = asyncio.run(create_dataset())
+    train(dataset)
+```
 
-#### 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]
+```
+@software{modern-bert-router,
+  author = {Reetz, Christian},
+  title = {ModernBERTRouter},
+  url = {https://huggingface.co/cdreetz/modern-bert-router/},
+  year = {2025}
+}
+```