Update README.md

README.md

@@ -18,23 +18,31 @@ Oneirogen can be used to generate novel dream narratives. It can also be used fo
 ## Code
 
  ```py
-from peft import prepare_model_for_kbit_training, PeftModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
 
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    )
+class CustomStoppingCriteria(StoppingCriteria):
+    def __init__(self, stop_token, tokenizer):
+        self.stop_token = stop_token
+        self.tokenizer = tokenizer
 
-model_id = "/gpfsdswork/projects/rech/uux/uvp47iv/notebooks/DreamLM/qwen2-7B-dream-final-merge"
+    def __call__(self, input_ids, scores, **kwargs):
+        decoded_output = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
+        if self.stop_token in decoded_output:
+            return True
+        return False
 
-#model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)#, attn_implementation="flash_attention_2")#, load_in_8bits=True)
+stop_token = "END." # The model was trained with this special end of text token.
+stopping_criteria = StoppingCriteriaList([CustomStoppingCriteria(stop_token, tokenizer)])
 
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, attn_implementation="flash_attention_2")#, load_in_8bits=True)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained("gustavecortal/oneirogen-0.5B")
+model = AutoModelForCausalLM.from_pretrained("gustavecortal/oneirogen-0.5B", torch_dtype=torch.float16)
+model.to("cuda")
+
+text = "Dream:" # The model was trained with this prefix
+
+inputs = tokenizer(text, return_tensors="pt").to("cuda")
+outputs = model.generate(inputs["input_ids"], attention_mask=inputs["attention_mask"], max_new_tokens=256, top_k = 50, top_p = 0.95, do_sample = True, temperature=0.9, num_beams = 1, repetition_penalty= 1.11, stopping_criteria=stopping_criteria)
+print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=False)[0])
 ```
 
 ## Inspiration
@@ -49,6 +57,4 @@ Mail: [email protected]
 
 X: [@gustavecortal](https://x.com/gustavecortal)
 
-Website: [gustavecortal.com](gustavecortal.com)
-
-
+Website: [gustavecortal.com](gustavecortal.com)