Update README.md

README.md

@@ -29,13 +29,13 @@ Mariana Moreira dos Santos ([LinkedIn](https://www.linkedin.com/in/mmoreirast/))
 You can check the codes used to fine-tune the model at the following [Google Colab](https://colab.research.google.com/drive/1SvJvTcH3IRnsEv72UxkVmV0oClCZARtE?usp=sharing) link. 
 
 ## Fine-tuning details
-- **Base model:** [TeenyTinyLlama 460m](https://huggingface.co/nicholasKluge/TeenyTinyLlama-460m)
+- **Base model:** [TeenyTinyLlama 160m](https://huggingface.co/nicholasKluge/TeenyTinyLlama-160m)
 - **Context length:** 2048 tokens
 - **Dataset for fine-tuning:** [medicine-training-pt](mmoreirast/medicine-training-pt)
 - **Dataset for evaluation:** [medicine-evaluation-pt](https://huggingface.co/datasets/mmoreirast/medicine-evaluation-pt)
 - **Language:** Portuguese
-- **GPU:** NVIDIA A100-SXM4-40GB
-- **Training time**: ~5 hours 
+- **GPU:** NVIDIA L4
+- **Training time**: ~9 hours 
 
 ## Parameters
 - **Number of Epochs:** 4
@@ -61,7 +61,7 @@ Using the `pipeline`:
 ```python
 from transformers import pipeline
 
-generator = pipeline("text-generation", model="mmoreirast/Doctor-Llama-460m")
+generator = pipeline("text-generation", model="mmoreirast/Doctor-Llama-160m")
 
 completions  = generator("Me fale sobre o sistema nervoso", num_return_sequences=2, max_new_tokens=100)
 
@@ -76,8 +76,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Load model and the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("mmoreirast/Doctor-Llama-460m", revision='main')
-model = AutoModelForCausalLM.from_pretrained("mmoreirast/Doctor-Llama-460m", revision='main')
+tokenizer = AutoTokenizer.from_pretrained("mmoreirast/Doctor-Llama-160m", revision='main')
+model = AutoModelForCausalLM.from_pretrained("mmoreirast/Doctor-Llama-160m", revision='main')
 
 # Pass the model to your device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")