Rename dataset generate to dataset.py

dataset generate

@@ -1,5 +0,0 @@
-from transformers import pipeline
-gen = pipeline("text2text-generation", model=MODEL, torch_dtype="auto", device_map="auto")
-
-def gen_one(ings):
-    return gen(f"ingredients: {ings}", max_new_tokens=180, do_sample=True, temperature=0.9, top_p=0.95)[0]["generated_text"]

dataset.py

@@ -0,0 +1,29 @@
+# make_dataset.py
+import random, json, torch
+from datasets import Dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+MODEL = "erwanlc/t5-cocktails_recipe-base"
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL, device_map="auto")
+
+ING_POOL = ["vodka","gin","rum","tequila","whiskey","triple sec","vermouth","lime juice",
+            "lemon juice","cranberry juice","pineapple juice","simple syrup","agave syrup",
+            "bitters","ginger beer","soda water","tonic water","mint","basil","cucumber"]
+
+def rand_ings():
+    return ", ".join(random.sample(ING_POOL, k=random.randint(3,6)))
+
+rows = []
+for i in range(1000):
+    ings = rand_ings()
+    prompt = f"ingredients: {ings}"
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    ids = model.generate(**inputs, max_new_tokens=180, do_sample=True, temperature=0.9, top_p=0.95)
+    text = tokenizer.decode(ids[0], skip_special_tokens=True)
+    rows.append({"id": i, "ingredients_text": ings, "generated_text": text})
+
+ds = Dataset.from_list(rows)
+ds.to_parquet("cocktail_synth.parquet")
+print("Wrote cocktail_synth.parquet")
+