Update README.md

README.md

@@ -15,49 +15,64 @@ should probably proofread and complete it, then remove this comment. -->
 
 # tool-bert
 
-This model is a fine-tuned version of [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased) on an unknown dataset.
-It achieves the following results on the evaluation set:
-- Loss: 0.0013
-- Accuracy: 1.0
+This model is a fine-tuned version of [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased).
 
-## Model description
+It uses a custom made dataset of sample user instructions, which are classified to a number of possible local assistant function calling endpoints. 
 
-More information needed
+For example, given an input query, tool-bert returns a prediction as to what tool to use to augment a downstream LLM generated output with. 
 
-## Intended uses & limitations
+More information on these tools to follow, but example tools are "play music", "check the weather", "get the news", "take a photo", or use no tool. 
 
-More information needed
+Basically, this model is meant to be a means of allowing very small LLMs (i.e. 8B and below) to use function calling. 
 
-## Training and evaluation data
+All limitations and biases are inherited from the parent model. 
 
-More information needed
+### Example Usage
 
-## Training procedure
+```python
+from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification
 
-### Training hyperparameters
+key_tools = ['take_picture', 'no_tool_needed',
+             'check_news', 'check_weather', 'play_spotify']
 
-The following hyperparameters were used during training:
-- learning_rate: 5e-05
-- train_batch_size: 8
-- eval_batch_size: 8
-- seed: 42
-- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
-- lr_scheduler_type: linear
-- num_epochs: 4
 
-### Training results
+def get_id2tool_name(id, key_tools):
+    return key_tools[id]
 
-| Training Loss | Epoch | Step | Validation Loss | Accuracy |
-|:-------------:|:-----:|:----:|:---------------:|:--------:|
-| No log        | 1.0   | 81   | 0.0077          | 1.0      |
-| No log        | 2.0   | 162  | 0.0022          | 1.0      |
-| No log        | 3.0   | 243  | 0.0014          | 1.0      |
-| No log        | 4.0   | 324  | 0.0013          | 1.0      |
 
+def remove_any_non_alphanumeric_characters(text):
+    return ''.join(e for e in text if e.isalnum() or e.isspace())
 
+
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "nkasmanoff/tool-bert")
+
+    model.eval()
+    return model, tokenizer
+
+
+def predict_tool(question, model, tokenizer):
+    question = remove_any_non_alphanumeric_characters(question)
+    inputs = tokenizer(question, return_tensors="pt")
+
+    outputs = model(**inputs)
+
+    logits = outputs.logits
+    return get_id2tool_name(logits.argmax().item(), key_tools)
+
+model, tokenizer = load_model()
+
+question = "What's the weather outside?"
+
+predict_tool(question, model, tokenizer)
+> check_weather
+```
 ### Framework versions
 
 - Transformers 4.41.1
 - Pytorch 2.3.0
 - Datasets 2.19.1
-- Tokenizers 0.19.1
+- Tokenizers 0.19.1