Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import torch | |
| from transformers import pipeline, set_seed | |
| from transformers import AutoTokenizer | |
| from PIL import ( | |
| Image, | |
| ImageFont, | |
| ImageDraw | |
| ) | |
| import re | |
| from examples import EXAMPLES | |
| import meta | |
| from utils import ( | |
| remote_css, | |
| local_css, | |
| load_image, | |
| pure_comma_separation | |
| ) | |
| class TextGeneration: | |
| def __init__(self): | |
| self.debug = True | |
| self.dummy_output = { | |
| 'directions': [ | |
| 'peel the potato and slice thinly.', | |
| 'place in a microwave safe dish.', | |
| 'cover with plastic wrap and microwave on high for 5 minutes.', | |
| 'remove from the microwave and sprinkle with cheese.', | |
| 'return to the microwave for 1 minute or until cheese is melted.', | |
| 'return to the microwave for 1 minute or until cheese is melted.', | |
| 'return to the microwave for 1 minute or until cheese is melted.' | |
| ], | |
| 'ingredients': [ | |
| '1 potato', | |
| '1 slice cheese', | |
| '1 potato', | |
| '1 slice cheese' | |
| '1 potato', | |
| '1 slice cheese', | |
| '1 slice cheese', | |
| '1 potato', | |
| '1 slice cheese' | |
| '1 potato', | |
| '1 slice cheese', | |
| ], | |
| 'title': 'Cheese Potatoes' | |
| } | |
| self.tokenizer = None | |
| self.generator = None | |
| self.task = "text2text-generation" | |
| self.model_name_or_path = "flax-community/t5-recipe-generation" | |
| self.h1_font = ImageFont.truetype("asset/fonts/PT_Serif/PTSerif-Bold.ttf", 100) | |
| self.h2_font = ImageFont.truetype("asset/fonts/PT_Serif/PTSerif-Bold.ttf", 50) | |
| self.p_font = ImageFont.truetype("asset/fonts/PT_Serif/PTSerif-Regular.ttf", 30) | |
| set_seed(42) | |
| def _skip_special_tokens_and_prettify(self, text): | |
| recipe_maps = {"<sep>": "--", "<section>": "\n"} | |
| recipe_map_pattern = "|".join(map(re.escape, recipe_maps.keys())) | |
| text = re.sub( | |
| recipe_map_pattern, | |
| lambda m: recipe_maps[m.group()], | |
| re.sub("|".join(self.tokenizer.all_special_tokens), "", text) | |
| ) | |
| data = {"title": "", "ingredients": [], "directions": []} | |
| for section in text.split("\n"): | |
| section = section.strip() | |
| if section.startswith("title:"): | |
| data["title"] = " ".join( | |
| [w.strip().capitalize() for w in section.replace("title:", "").strip().split() if w.strip()] | |
| ) | |
| elif section.startswith("ingredients:"): | |
| data["ingredients"] = [s.strip() for s in section.replace("ingredients:", "").split('--')] | |
| elif section.startswith("directions:"): | |
| data["directions"] = [s.strip() for s in section.replace("directions:", "").split('--')] | |
| else: | |
| pass | |
| return data | |
| def load(self): | |
| if not self.debug: | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path) | |
| self.generator = pipeline(self.task, model=self.model_name_or_path, tokenizer=self.model_name_or_path) | |
| def prepare_frame(self, recipe, frame): | |
| im_editable = ImageDraw.Draw(frame) | |
| # Title | |
| ws, hs = 120, 500 | |
| im_editable.text( | |
| (ws, hs), | |
| recipe["title"], | |
| (61, 61, 70), | |
| font=self.h1_font, | |
| ) | |
| # Ingredients | |
| hs = hs + 200 | |
| im_editable.text( | |
| (ws, hs), | |
| "Ingredients", | |
| (61, 61, 70), | |
| font=self.h2_font, | |
| ) | |
| hs = hs + 80 | |
| im_editable.text( | |
| (ws + 10, hs), | |
| "\n".join([f"- {item}" for item in recipe["ingredients"]]), | |
| (61, 61, 70), | |
| font=self.p_font, | |
| ) | |
| # Directions | |
| hs = hs + 400 | |
| im_editable.text( | |
| (ws, hs), | |
| "Directions", | |
| (61, 61, 70), | |
| font=self.h2_font, | |
| ) | |
| hs = hs + 80 | |
| im_editable.text( | |
| (ws + 10, hs), | |
| "\n".join([f"- {item}" for item in recipe["directions"]]), | |
| (61, 61, 70), | |
| font=self.p_font, | |
| ) | |
| return frame | |
| def generate(self, items, generation_kwargs): | |
| print(generation_kwargs) | |
| if not self.debug: | |
| generation_kwargs["num_return_sequences"] = 1 | |
| # generation_kwargs["return_full_text"] = False | |
| generation_kwargs["return_tensors"] = True | |
| generation_kwargs["return_text"] = False | |
| generated_ids = self.generator( | |
| items, | |
| **generation_kwargs, | |
| )[0]["generated_token_ids"] | |
| recipe = self.tokenizer.decode(generated_ids, skip_special_tokens=False) | |
| recipe = self._skip_special_tokens_and_prettify(recipe) | |
| return recipe | |
| return self.dummy_output | |
| def generate_frame(self, recipe): | |
| frame = load_image("asset/images/recipe-post.png") | |
| return self.prepare_frame(recipe, frame) | |
| def load_text_generator(): | |
| generator = TextGeneration() | |
| generator.load() | |
| return generator | |
| chef_top = { | |
| "max_length": 512, | |
| "min_length": 64, | |
| "no_repeat_ngram_size": 3, | |
| "do_sample": True, | |
| "top_k": 60, | |
| "top_p": 0.95, | |
| "num_return_sequences": 1 | |
| } | |
| chef_beam = { | |
| "max_length": 512, | |
| "min_length": 64, | |
| "no_repeat_ngram_size": 3, | |
| "early_stopping": True, | |
| "num_beams": 5, | |
| "length_penalty": 1.5, | |
| "num_return_sequences": 1 | |
| } | |
| def main(): | |
| st.set_page_config( | |
| page_title="Chef Transformer", | |
| page_icon="🍲", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| generator = load_text_generator() | |
| local_css("asset/css/style.css") | |
| st.sidebar.image(load_image("asset/images/chef-transformer-transparent.png"), width=310) | |
| st.sidebar.title("Choose your own chef") | |
| chef = st.sidebar.selectbox("Chef", index=0, options=["Chef Scheherazade", "Chef Giovanni"]) | |
| st.markdown(meta.HEADER_INFO) | |
| prompts = list(EXAMPLES.keys()) + ["Custom"] | |
| prompt = st.selectbox('Examples', prompts, index=len(prompts) - 1) | |
| if prompt == "Custom": | |
| prompt_box = "" | |
| else: | |
| prompt_box = EXAMPLES[prompt] | |
| items = st.text_input( | |
| 'Add custom ingredients here (separated by `,`): ', | |
| pure_comma_separation(prompt_box, return_list=False), | |
| key="custom_keywords", | |
| max_chars=1000) | |
| items = pure_comma_separation(items, return_list=False) | |
| entered_items = st.empty() | |
| if st.button('Get Recipe!'): | |
| entered_items.markdown("**Generate recipe for:** " + items) | |
| with st.spinner("Generating recipe..."): | |
| gen_kw = chef_top if chef == "Chef Scheherazade" else chef_beam | |
| generated_recipe = generator.generate(items, gen_kw) | |
| recipe_post = generator.generate_frame(generated_recipe) | |
| col1, col2, col3 = st.beta_columns([1, 6, 1]) | |
| with col1: | |
| st.write("") | |
| with col2: | |
| st.image( | |
| recipe_post, | |
| # width=500, | |
| caption="Your recipe", | |
| use_column_width="auto", | |
| output_format="PNG" | |
| ) | |
| with col3: | |
| st.write("") | |
| if __name__ == '__main__': | |
| main() | |