replace gpt2 by Mistral

update version to 0.1.3

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,

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:

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:

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:

transformers_cfg/__init__.py

@@ -2,4 +2,4 @@
 
 setup_logging()
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"