Update README.md

README.md

@@ -23,10 +23,17 @@ It achieves the following results on the evaluation set:
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline
+from peft import PeftConfig, PeftModel
+import torch
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map
 
-base_model = "CLARA-MeD/bertin-gpt"
-tokenizer = AutoTokenizer.from_pretrained(base_model)
-model = AutoModelForCausalLM.from_pretrained(base_model).cuda()
+
+repo_name = "CLARA-MeD/bertin-gpt"
+config = PeftConfig.from_pretrained(repo_name)
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path,torch_dtype=torch.float16,
+    device_map="auto")
+model = PeftModel.from_pretrained(model, repo_name)
 ```
 
 For generation, we can use the model's `.generate()` method. Remember that the prompt needs a **Spanish** template:
@@ -57,7 +64,9 @@ Simplifica la siguiente frase
         output = tokenizer.decode(seq, skip_special_tokens=True)
         print(output.split("### Respuesta:")[-1].strip())
 
-generate("Al sujeto se le ha tratado previamente con antagonistas del factor de necrosis tumoral alfa (TNF-α) sin respuesta clínica documentada al tratamiento. También puede ocurrir que al sujeto no se le tratara anteriormente con antagonistas de TNF-α, pero está fallando el tratamiento convencional actual.")
+generate("Acromegalia")
+# La acromegalia es un trastorno causado por un exceso de hormona del crecimiento en el cuerpo.
+
 
 ```