test

generate.py

@@ -20,12 +20,14 @@ logger = logging.getLogger(__name__)
 
 
 logger.warning("Loading model...")
-model_id = "google/gemma-2b-it"
-# model_id = "Qwen/Qwen1.5-0.5B-Chat"
+# model_id = "google/gemma-2b-it"
+model_id = "Qwen/Qwen1.5-0.5B-Chat"
 if torch.backends.mps.is_available():
-    model = models.transformers(model_id, device="mps")
+    device = "mps"
+    model = models.transformers(model_id, device=device)
 else:
-    model = models.transformers(model_id, device="cuda")
+    device = "cuda"
+    model = models.transformers(model_id, device=device)
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 sampler = PenalizedMultinomialSampler()
@@ -69,7 +71,7 @@ def get_samples_generator(new_fields: list[str]) -> SequenceGenerator:
         fsm=fsm,
         model=samples_generator_template.model,
         sampler=samples_generator_template.sampler,
-        device=samples_generator_template.device
+        device=device
     )