| from transformers import FlaxAutoModelForSeq2SeqLM |
| from transformers import AutoTokenizer |
| import textwrap |
|
|
| MODEL_NAME_OR_PATH = "JustAPR/resGen" |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH, use_fast=True) |
| model = FlaxAutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME_OR_PATH) |
|
|
| prefix = "items: " |
| |
| |
| |
| |
| |
| |
| |
| |
| generation_kwargs = { |
| "max_length": 512, |
| "min_length": 64, |
| "no_repeat_ngram_size": 3, |
| "early_stopping": True, |
| "num_beams": 5, |
| "length_penalty": 1.5, |
| } |
|
|
| special_tokens = tokenizer.all_special_tokens |
| tokens_map = { |
| "<sep>": "--", |
| "<section>": "\n" |
| } |
| def skip_special_tokens(text, special_tokens): |
| for token in special_tokens: |
| text = text.replace(token, '') |
|
|
| return text |
|
|
| def target_postprocessing(texts, special_tokens): |
| if not isinstance(texts, list): |
| texts = [texts] |
| |
| new_texts = [] |
| for text in texts: |
| text = skip_special_tokens(text, special_tokens) |
|
|
| for k, v in tokens_map.items(): |
| text = text.replace(k, v) |
|
|
| new_texts.append(text) |
|
|
| return new_texts |
|
|
| def generation_function(texts): |
| _inputs = texts if isinstance(texts, list) else [texts] |
| inputs = [prefix + inp for inp in _inputs] |
| inputs = tokenizer( |
| inputs, |
| max_length=256, |
| padding="max_length", |
| truncation=True, |
| return_tensors='jax' |
| ) |
|
|
| input_ids = inputs.input_ids |
| attention_mask = inputs.attention_mask |
|
|
| output_ids = model.generate( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| **generation_kwargs |
| ) |
| generated = output_ids.sequences |
| generated_recipe = target_postprocessing( |
| tokenizer.batch_decode(generated, skip_special_tokens=False), |
| special_tokens |
| ) |
| return generated_recipe |
|
|
|
|
| items = [ |
| "macaroni, butter, salt, bacon, milk, flour, pepper, cream corn", |
| "provolone cheese, bacon, bread, ginger" |
| ] |
| generated = generation_function(items) |
| for text in generated: |
| sections = text.split("\n") |
| for section in sections: |
| section = section.strip() |
| if section.startswith("title:"): |
| section = section.replace("title:", "") |
| headline = "TITLE" |
| elif section.startswith("ingredients:"): |
| section = section.replace("ingredients:", "") |
| headline = "INGREDIENTS" |
| elif section.startswith("directions:"): |
| section = section.replace("directions:", "") |
| headline = "DIRECTIONS" |
| |
| if headline == "TITLE": |
| print(f"[{headline}]: {section.strip().capitalize()}") |
| else: |
| section_info = [f" - {i+1}: {info.strip().capitalize()}" for i, info in enumerate(section.split("--"))] |
| print(f"[{headline}]:") |
| print("\n".join(section_info)) |
|
|
| print("-" * 130) |
