Update README.md

README.md

@@ -93,7 +93,7 @@ active_params/tok : ~160,000
 ## Usage
 
 ```python
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 from tokenizer import BPETokenizer
 import torch
 
@@ -103,9 +103,25 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-# model.vocab and model.model must be in the same folder
 tok = BPETokenizer.load("model.vocab", "model.model")
 
+# Bloqueia EOS e PAD nos primeiros min_tokens gerados
+class MinNewTokens(LogitsProcessor):
+    def __init__(self, min_tokens: int, eos_id: int, pad_id: int):
+        self.min_tokens = min_tokens
+        self.bad = [eos_id, pad_id]
+        self.generated = 0
+
+    def __call__(self, input_ids, scores):
+        if self.generated < self.min_tokens:
+            for bid in self.bad:
+                scores[:, bid] = float("-inf")
+        self.generated += 1
+        return scores
+
+eos_id = tok.token2id["<eos>"]
+pad_id = tok.token2id["<pad>"]
+
 prompt = "# Pergunta:\nQuanto é 5 + 3?\n--\n# Resposta:\n"
 ids = tok.encode(prompt, add_bos=True, add_eos=False)
 input_ids = torch.tensor([ids])
@@ -113,16 +129,25 @@ input_ids = torch.tensor([ids])
 with torch.no_grad():
     output = model.generate(
         input_ids,
-        max_new_tokens=60,
-        temperature=0.8,
+        max_new_tokens=80,
+        temperature=0.9,
         do_sample=True,
-        top_k=40,
-        top_p=0.9,
-        eos_token_id=tok.token2id["<eos>"],
-        pad_token_id=tok.token2id["<pad>"],
+        top_k=50,
+        top_p=0.95,
+        eos_token_id=eos_id,
+        pad_token_id=pad_id,
+        use_cache=False,
+        logits_processor=LogitsProcessorList([
+            MinNewTokens(min_tokens=5, eos_id=eos_id, pad_id=pad_id)
+        ]),
     )
 
-print(tok.decode(output[0].tolist()))
+new_tokens = output[0][len(ids):].tolist()
+# Remove EOS do final se presente
+if new_tokens and new_tokens[-1] == eos_id:
+    new_tokens = new_tokens[:-1]
+
+print("Resposta:", tok.decode(new_tokens))
 ```
 
 ---