From 1f3e9875495c73b1d18606704ac5c80f31a5e7b0 Mon Sep 17 00:00:00 2001
From: Saibo Geng <saibo.geng@yahoo.fr>
Date: Tue, 27 Feb 2024 20:14:03 +0100
Subject: [PATCH] replace gpt2 by Mistral update version to 0.1.3

---
 README.md                       | 18 +++++++++++++-----
 examples/generate_cIE.py        |  8 ++++++--
 examples/generate_json.py       |  8 ++++++--
 examples/generate_json_array.py |  8 ++++++--
 transformers_cfg/__init__.py    |  2 +-
 5 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 874ff6a..0d7087f 100644
--- a/README.md
+++ b/README.md
@@ -30,16 +30,26 @@ pip install git+https://github.com/epfl-dlab/transformers-CFG.git
 The below example can be found in `examples/generate_json.py`
 
 ```python
-
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers_cfg.grammar_utils import IncrementalGrammarConstraint
 from transformers_cfg.generation.logits_process import GrammarConstrainedLogitsProcessor
 
 if __name__ == "__main__":
+    # Detect if GPU is available, otherwise use CPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    model_id = "mistralai/Mistral-7B-v0.1"
+
     # Load model and tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+    model = AutoModelForCausalLM.from_pretrained(model_id).to(
+        device
+    )  # Load model to defined device
+    model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
     # Load json grammar
     with open("examples/grammars/json.ebnf", "r") as file:
@@ -54,9 +64,7 @@ if __name__ == "__main__":
 
     output = model.generate(
         input_ids,
-        do_sample=False,
         max_length=50,
-        num_beams=2,
         logits_processor=[grammar_processor],
         repetition_penalty=1.1,
         num_return_sequences=1,
diff --git a/examples/generate_cIE.py b/examples/generate_cIE.py
index 88c2e41..f889f21 100644
--- a/examples/generate_cIE.py
+++ b/examples/generate_cIE.py
@@ -10,12 +10,16 @@
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
 
+    model_id = "mistralai/Mistral-7B-v0.1"
+
     # Load model and tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained("gpt2").to(
+
+    model = AutoModelForCausalLM.from_pretrained(model_id).to(
         device
     )  # Load model to defined device
+    model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
     # Load grammar
     with open("examples/grammars/cIE.ebnf", "r") as file:
diff --git a/examples/generate_json.py b/examples/generate_json.py
index caa7192..1f0fb4b 100644
--- a/examples/generate_json.py
+++ b/examples/generate_json.py
@@ -10,12 +10,16 @@
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
 
+    model_id = "mistralai/Mistral-7B-v0.1"
+
     # Load model and tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained("gpt2").to(
+
+    model = AutoModelForCausalLM.from_pretrained(model_id).to(
         device
     )  # Load model to defined device
+    model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
     # Load grammar
     with open("examples/grammars/json.ebnf", "r") as file:
diff --git a/examples/generate_json_array.py b/examples/generate_json_array.py
index 0f037ac..5d7b7b7 100644
--- a/examples/generate_json_array.py
+++ b/examples/generate_json_array.py
@@ -9,12 +9,16 @@
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
 
+    model_id = "mistralai/Mistral-7B-v0.1"
+
     # Load model and tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained("gpt2").to(
+
+    model = AutoModelForCausalLM.from_pretrained(model_id).to(
         device
     )  # Load model to defined device
+    model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
     # Load grammar
     with open("examples/grammars/json_arr.ebnf", "r") as file:
diff --git a/transformers_cfg/__init__.py b/transformers_cfg/__init__.py
index e67e5c7..ea02d19 100644
--- a/transformers_cfg/__init__.py
+++ b/transformers_cfg/__init__.py
@@ -2,4 +2,4 @@
 
 setup_logging()
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"