Update README.md

README.md

@@ -80,10 +80,11 @@ To use this model for inference, you need to load the fine-tuned model and token
 
 Running on CPU
 ```python
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-tokenizer = T5Tokenizer.from_pretrained('Falconsai/arc_of_conversation')
-model = T5ForConditionalGeneration.from_pretrained('Falconsai/arc_of_conversation')
+tokenizer = AutoTokenizer.from_pretrained("Falconsai/arc_of_conversation")
+model = AutoModelForSeq2SeqLM.from_pretrained("Falconsai/arc_of_conversation")
 
 input_text = "Your conversation Here"
 input_ids = tokenizer(input_text, return_tensors="pt").input_ids
@@ -95,10 +96,10 @@ print(tokenizer.decode(outputs[0]))
 Running on GPU
 ```python
 # pip install accelerate
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-tokenizer = T5Tokenizer.from_pretrained('Falconsai/arc_of_conversation')
-model = T5ForConditionalGeneration.from_pretrained('Falconsai/arc_of_conversation', device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("Falconsai/arc_of_conversation")
+model = AutoModelForSeq2SeqLM.from_pretrained("Falconsai/arc_of_conversation", device_map="auto")
 
 input_text = "Your conversation Here"
 input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")