Fix ort documentation code snippet (#2070)

fix code snippet
This commit is contained in:
Ella Charlaix
2024-10-19 20:37:48 +02:00
committed by GitHub
parent 59d6f7f04e
commit 8af46e53bd
@@ -13,11 +13,11 @@ Once your model was [exported to the ONNX format](https://huggingface.co/docs/op
```diff
from transformers import AutoTokenizer, pipeline
- from transformers import AutoModelForQuestionAnswering
+ from optimum.onnxruntime import ORTModelForQuestionAnswering
- from transformers import AutoModelForCausalLM
+ from optimum.onnxruntime import ORTModelForCausalLM
- model = AutoModelForQuestionAnswering.from_pretrained("meta-llama/Llama-3.2-1B) # PyTorch checkpoint
+ model = ORTModelForQuestionAnswering.from_pretrained("onnx-community/Llama-3.2-1B", subfolder="onnx") # ONNX checkpoint
- model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B) # PyTorch checkpoint
+ model = ORTModelForCausalLM.from_pretrained("onnx-community/Llama-3.2-1B", subfolder="onnx") # ONNX checkpoint
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
@@ -74,4 +74,4 @@ You can also call `push_to_hub` directly on your model to upload it to the [Hub]
# Push the onnx model to HF Hub
>>> model.push_to_hub(output_dir, repository_id="my-onnx-repo") # doctest: +SKIP
```
```