lower columns temperature

generate.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import time
+from pathlib import Path
 from typing import Annotated, Iterator
 
 import ijson
@@ -31,6 +32,7 @@ else:
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 sampler = PenalizedMultinomialSampler()
+low_temperature_sampler = PenalizedMultinomialSampler(temperature=0.3)
 empty_tokens = [token_id for token_id in range(tokenizer.vocab_size) if not tokenizer.decode([token_id]).strip()]
 sampler.set_max_repeats(empty_tokens, 1)
 
@@ -56,7 +58,7 @@ samples_generator_template = generate.json(model, Dataset, sampler=sampler)
 class Columns(BaseModel):
     columns: conset(Annotated[str, StringConstraints(pattern=r'[a-z0-9_]+')], min_length=2, max_length=len(Sample.model_fields) - 1)  # type: ignore
 
-columns_generator = generate.json(model, Columns, sampler=sampler)
+columns_generator = generate.json(model, Columns, sampler=low_temperature_sampler)
 
 def get_samples_generator(new_fields: list[str]) -> SequenceGenerator:
     fsm=samples_generator_template.fsm
@@ -89,7 +91,8 @@ def samples_prommpt(filename: str, prompt: str, columns: str):
     {{ prompt }}
     """
 
-def stream_file(filename: str, prompt: str, columns: list[str], seed: int, size: int) -> Iterator[str]:
+def stream_jsonl_file(filename: str, prompt: str, columns: list[str], seed: int, size: int) -> Iterator[str]:
+    filename = Path(filename).stem
     logger.warning(f"stream_response({filename=}, {prompt=}, {columns=})")
     _start = time.time()
     rng = torch.Generator(device=model.device)

samplers.py

@@ -6,6 +6,7 @@ from outlines.samplers import MultinomialSampler
 
 logger = logging.getLogger(__name__)
 
+
 class PenalizedMultinomialSampler(MultinomialSampler):
 
     def __init__(self, **kwargs):