Update README.md

README.md

@@ -59,19 +59,21 @@ training_args = Seq2SeqTrainingArguments(
 
 ```
 import torch
-from transformers import AutoTokenizer, T5ForConditionalGeneration
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the tokenizer and model
 model_path = 'juanfra218/text2sql'
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+tokenizer = T5Tokenizer.from_pretrained(model_path)
 model = T5ForConditionalGeneration.from_pretrained(model_path)
+model.to(device)  
 
 # Function to generate SQL queries
 def generate_sql(prompt, schema):
     input_text = "translate English to SQL: " + prompt + " " + schema
     inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True, padding="max_length")
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     inputs = {key: value.to(device) for key, value in inputs.items()}
 
     max_output_length = 1024