Update README.md

README.md

@@ -82,22 +82,17 @@ Additionally to conventional generation, using Tool-augmented generation require
 (1) implementation of the tool(s) and 
 (2) a customization of generate() method augmenting input context on-demand with the outputs of the tools.
 
-You can find these two components implemented in the **gadgets/gadget_assisted_model.py** and **gadgets/gadget.py** in the project's [home repo](https://github.com/emnlp2023sub/gadgets).
+You can find these two components implemented in the **gadgets/models.py** and **gadgets/gadget.py** in the project's [home repo](https://github.com/emnlp2023sub/gadgets).
 
 After adding these two scripts to your directory, you can use the model as follows:
 
 ```python
-from gadget_assisted_model import GadgetAssistedModel
-from gadget import Calculator
-
 from transformers import T5ForConditionalGeneration, T5Tokenizer
 
+from gadgets.model import gadget_assisted_model
+from gadgets.gadget import Calculator
 
-class GadgetAssistedT5(GadgetAssistedModel, T5ForConditionalGeneration):
-    # GadgetAssistedModel overrides the standard generate() from transformers
-    pass
-
-
+GadgetAssistedT5 = gadget_assisted_model(T5ForConditionalGeneration)
 model = GadgetAssistedT5.from_pretrained("emnlp2023/calc-t5-xl")
 tokenizer = T5Tokenizer.from_pretrained("emnlp2023/calc-t5-xl")
 
@@ -117,8 +112,10 @@ tokenizer.decode(output_ids[0], spaces_between_special_tokens=False)
 ```
 This returns:
 ```html
-According to the ratio, Mike got 2/5*$2500 = $<gadget id="calculator">2/5*2500</gadget><output>1_000</output> 1000 
-Mike will have $1000-$200 = $<gadget id="calculator">1000-200</gadget><output>800</output> 800 after buying a shirt. 
+According to the ratio, for every 5 parts that Johnson gets, Mike gets 2 parts Since Johnson got $2500,
+each part is therefore $2500/5 = $<gadget id="calculator">2500/5</gadget><output>500</output> 500
+Mike will get 2*$500 = $<gadget id="calculator">2*500</gadget><output>1_000</output> 1000
+After buying the shirt he will have $1000-$200 = $<gadget id="calculator">1000-200</gadget><output>800</output> 800 left.
 Final result is<result>800</result></s>
 ```