[bug] transformers generator memory monster

[cdreetz]

need to reset model or something after each generation i think or clear cache

import os
import chatan
from transformers import AutoTokenizer, AutoModelForCausalLM

#model = AutoModelForCausalLM.from_pretrained("cdreetz/kwen2.5-1.5b-v2")
#tokenizer = AutoTokenizer.from_pretrained("cdreetz/kwen2.5-1.5b-v2")
#
#prompt = "Write a Triton kernel for element-wise addition:"
#inputs = tokenizer(prompt, return_tensors="pt")
#outputs = model.generate(**inputs, max_new_tokens=1024)
#response = tokenizer.decode(outputs[0], skip_special_tokens=True)
#print(response)


gen = chatan.generator("transformers", model="cdreetz/kwen2.5-1.5b-v2")
judge = chatan.generator("openai", os.getenv("OPENAI_API_KEY"), model="o3")

ds = chatan.dataset({
    "prompt": "Write a Triton kernel for element-wise addition",
    "response": gen("{prompt}"),
    "score": judge("is this kernel implementation correct?  {response} respond with 2 if yes, 1 if almost, 0 if no. only respond with the score.")
}, n=2)

out = ds.generate()

print(out)

christian@Mac try-kwen % uv run kwen-chatan.py
Loading checkpoint shards: 100%|████████████████████████████████| 2/2 [00:00<00:00, 32.70it/s]
Device set to use mps:0
Generating:   0%|                                                       | 0/2 [00:00<?, ?it/s]/AppleInternal/Library/BuildRoots/5a8a3fcc-55cb-11ef-848e-8a553ba56670/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:850: failed assertion `[MPSTemporaryNDArray initWithDevice:descriptor:isTextureBacked:] Error: total bytes of NDArray > 2**32'
/Users/christian/.local/share/uv/python/cpython-3.11.11-macos-aarch64-none/lib/python3.11/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '