Update README.md

README.md

@@ -34,12 +34,17 @@ Gemma-2B Fine-Tuned Python Model is a deep learning model based on the Gemma-2B
 1. **Install Gemma Python Package**:
    ```bash
     pip install -q -U transformers==4.38.0
+    pip install torch
    ```
 
 ## Inference
+
 1. **How to use the model in our notebook**:
 ```python
+
+
 # Load model directly
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("suriya7/Gemma-2B-Finetuned-Python-Model")
@@ -51,12 +56,15 @@ prompt_template = f"""
 <end_of_turn>\n<start_of_turn>model
 """
 prompt = prompt_template
-encodeds = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+encodeds = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).input_ids
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model.to(device)
+inputs = encodeds.to(device)
 
-model_inputs = encodeds.to('cuda')
 
 # Increase max_new_tokens if needed
-generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=False, pad_token_id=tokenizer.eos_token_id)
+generated_ids = model.generate(inputs, max_new_tokens=1000, do_sample=False, pad_token_id=tokenizer.eos_token_id)
 ans = ''
 for i in tokenizer.decode(generated_ids[0], skip_special_tokens=True).split('<end_of_turn>')[:2]:
     ans += i