Update README.md

README.md

@@ -1,8 +1,8 @@
 ---
 license: apache-2.0
 datasets:
-- mmoreirast/medicine-training-pt
 - mmoreirast/medicine-evaluation-pt
+- mmoreirast/aira-med-training-pt
 language:
 - pt
 metrics:
@@ -31,7 +31,7 @@ You can check the codes used to fine-tune the model at the following [Google Col
 ## Fine-tuning details
 - **Base model:** [TeenyTinyLlama 460m](https://huggingface.co/nicholasKluge/TeenyTinyLlama-460m)
 - **Context length:** 2048 tokens
-- **Dataset for fine-tuning:** [medicine-training-pt](mmoreirast/medicine-training-pt)
+- **Dataset for fine-tuning:** [aira-med-training-pt](https://huggingface.co/datasets/mmoreirast/aira-med-training-pt)
 - **Dataset for evaluation:** [medicine-evaluation-pt](https://huggingface.co/datasets/mmoreirast/medicine-evaluation-pt)
 - **Language:** Portuguese
 - **GPU:** NVIDIA A100-SXM4-40GB
@@ -39,7 +39,7 @@ You can check the codes used to fine-tune the model at the following [Google Col
 
 ## Parameters
 - **Number of Epochs:** 4
-- **Batch size:** 8
+- **Batch size:** 3
 - **Optimizer:** torch.optim.AdamW (warmup_steps = 1e3, learning_rate = 1e-5, epsilon = 1e-8)
 
 ## Evaluations
@@ -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-Chat")
 
 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-Chat", revision='main')
+model = AutoModelForCausalLM.from_pretrained("mmoreirast/Doctor-Llama-Chat", revision='main')
 
 # Pass the model to your device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")