app.py

import streamlit as st
import re,torch
import json,os
from transformers import AutoModelForCausalLM, AutoTokenizer
from datetime import datetime
from huggingface_hub import login

login(token=os.getenv("TOKEN"))

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    torch_dtype="auto",
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
if 'recipe' not in st.session_state:
    st.session_state.recipe = None

if 'recipe_saved' not in st.session_state:
    st.session_state.recipe_saved = None

if 'user_direction' not in st.session_state:
    st.session_state.user_direction = None

if 'serving_size' not in st.session_state:
    st.session_state.serving_size = 2

if 'selected_difficulty' not in st.session_state:
    st.session_state.selected_difficulty = "Quick & Easy"

if 'exclusions' not in st.session_state:
    st.session_state.exclusions = None
    
def create_detailed_prompt(user_direction, exclusions, serving_size, difficulty):
    if difficulty == "Quick & Easy":
        prompt = (
            f"Provide a 'Quick and Easy' recipe for {user_direction} that excludes {exclusions} and has a serving size of {serving_size}. "
            f"It should require as few ingredients as possible and should be ready in as little time as possible. "
            f"The steps should be simple, and the ingredients should be commonly found in a household pantry. "
            f"Provide a detailed ingredient list and step-by-step guide that explains the instructions to prepare in detail."
        )
    elif difficulty == "Intermediate":
        prompt = (
            f"Provide a classic recipe for {user_direction} that excludes {exclusions} and has a serving size of {serving_size}. "
            f"The recipe should offer a bit of a cooking challenge but should not require professional skills. "
            f"The recipe should feature traditional ingredients and techniques that are authentic to its cuisine. "
            f"Provide a detailed ingredient list and step-by-step guide that explains the instructions to prepare in detail."
        )
    elif difficulty == "Professional":
        prompt = (
            f"Provide a advanced recipe for {user_direction} that excludes {exclusions} and has a serving size of {serving_size}. "
            f"The recipe should push the boundaries of culinary arts, integrating unique ingredients, advanced cooking techniques, and innovative presentations. "
            f"The recipe should be able to be served at a high-end restaurant or would impress at a gourmet food competition. "
            f"Provide a detailed ingredient list and step-by-step guide that explains the instructions to prepare in detail."
        )
    return prompt

def generate_recipe(user_inputs):
    with st.spinner('Building the perfect recipe...'):
        # provide_recipe_schema = {
        #     'type': 'function',
        #     'function': {
        #         'name': 'provide_recipe',
        #         'description': 'Provides a detailed recipe strictly adhering to the user input/specifications, especially ingredient exclusions and the recipe difficulty',
        #         'parameters': {
        #             'type': 'object',
        #             'properties': {
        #                 'name': {
        #                     'type': 'string',
        #                     'description': 'A creative name for the recipe'
        #                 },
        #                 'description': {
        #                     'type': 'string',
        #                     'description': 'a brief one-sentence description of the provided recipe'
        #                 },
        #                 'ingredients': {
        #                     'type': 'array',
        #                     'items': {
        #                         'type': 'object',
        #                         'properties': {
        #                             'name': {
        #                                 'type': 'string',
        #                                 'description': 'Quantity and name of the ingredient'
        #                             }
        #                         }
        #                     }
        #                 },
        #                 'instructions': {
        #                     'type': 'array',
        #                     'items': {
        #                         'type': 'object',
        #                         'properties': {
        #                             'step_number': {
        #                                 'type': 'number',
        #                                 'description': 'The sequence number of this step'
        #                             },
        #                             'instruction': {
        #                                 'type': 'string',
        #                                 'description': 'Detailed description of what to do in this step'
        #                             }
        #                         }
        #                     }
        #                 }
        #             },
        #             'required': [
        #                 'name',
        #                 'description',
        #                 'ingredients',
        #                 'instructions'
        #             ]
        #         }
        #     }
        # }
        prompt = create_detailed_prompt(user_inputs['user_direction'], user_inputs['exclusions'], user_inputs['serving_size'], user_inputs['difficulty'])
        # messages = [{"role": "user", "content": prompt}]
        # tool_section = "\n".join([f"{tool['function']['name']}({json.dumps(tool['function']['parameters'])})" for tool in [provide_recipe_schema]])
        # text = f"{prompt}\n\nTools:\n{tool_section}"
        
        # Tokenize and move to the correct device
        model_inputs = tokenizer(prompt, return_tensors="pt")
        # text = tokenizer.apply_chat_template(
        #     messages,
        #     tokenize=False,
        #     add_generation_prompt=True,
        #     tools=[provide_recipe_schema]
        # )

        # Tokenize and move to the correct device
        # model_inputs = tokenizer([text], return_tensors="pt")
        # torch.cuda.empty_cache()
        # with torch.no_grad():
        generated_ids = model.generate(
            **model_inputs,
            # max_new_tokens=512,
        )

        # generated_ids = [
        #     output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        # ]

        st.session_state.recipe = tokenizer.decode(generated_ids[0])
        st.session_state.recipe_saved = False
        
def clear_inputs():
    st.session_state.user_direction = None
    st.session_state.exclusions = None
    st.session_state.serving_size = 2
    st.session_state.selected_difficulty = "Quick & Easy"
    
st.title("Let's get cooking")
st.session_state.user_direction = st.text_area(
    "What do you want to cook? Describe anything - a dish, cuisine, event, or vibe.",
    value = st.session_state.user_direction,
    placeholder="quick snack, asian style bowl with either noodles or rice, something italian",
    )

st.session_state.serving_size = st.number_input(
    "How many servings would you like to cook?",
    min_value=1,
    max_value=100,
    value=st.session_state.serving_size,
    step=1
)

difficulty_dictionary = {
    "Quick & Easy": {
        "description": "Easy recipes with straightforward instructions. Ideal for beginners or those seeking quick and simple cooking.",
    },
    "Intermediate": {
        "description": "Recipes with some intricate steps that invite a little challenge. Perfect for regular cooks wanting to expand their repertoire with new ingredients and techniques.",
    },
    "Professional": {
        "description": "Complex recipes that demand a high level of skill and precision. Suited for seasoned cooks aspiring to professional-level sophistication and creativity.",
    }
}

st.session_state.selected_difficulty = st.radio(
    "Choose a difficulty level for your recipe.",
    [
        list(difficulty_dictionary.keys())[0], 
        list(difficulty_dictionary.keys())[1], 
        list(difficulty_dictionary.keys())[2]
    ],
    captions = [
        difficulty_dictionary["Quick & Easy"]["description"], 
        difficulty_dictionary["Intermediate"]["description"],
        difficulty_dictionary["Professional"]["description"]
    ],
    index=list(difficulty_dictionary).index(st.session_state.selected_difficulty)
)

st.session_state.exclusions = st.text_area(
    "Any ingredients you want to exclude?",
    value = st.session_state.exclusions,
    placeholder="gluten, dairy, nuts, cilantro",
    )

fancy_exclusions =""

if st.session_state.selected_difficulty == "Professional":
    exclude_fancy = st.checkbox(
        "Exclude cliche professional ingredients? (gold leaf, truffle, edible flowers, microgreens)", 
        value=True)
    fancy_exclusions = "gold leaf, truffle, edible flowers, microgreens, gold dust"
      

user_inputs = {
    "user_direction" : st.session_state.user_direction,
    "exclusions": f"{st.session_state.exclusions}, {fancy_exclusions}",
    "serving_size": st.session_state.serving_size,
    "difficulty": st.session_state.selected_difficulty
}
button_cols_submit = st.columns([1, 1, 4])
with button_cols_submit[0]:
    st.button(label='Submit', on_click=generate_recipe, kwargs=dict(user_inputs=user_inputs), type="primary", use_container_width=True)
with button_cols_submit[1]:
    st.button(label='Reset', on_click=clear_inputs, type="secondary", use_container_width=True)
with button_cols_submit[2]:
    st.empty()
    
if st.session_state.recipe is not None:
    st.divider()
    print(st.session_state.recipe)
    recipe = json.loads(st.session_state.recipe)
    recipe_md = ''
    recipe_md += f'# {recipe["name"]} \n\n'
    recipe_md += f'{recipe["description"]} \n\n'
    recipe_md += '## Ingredients: \n'
    for ingredient in recipe['ingredients']:
        recipe_md += f"- {ingredient['name']} \n"
    recipe_md += '\n## Instructions:\n'
    for instruction in recipe['instructions']:
        recipe_md += f"{instruction['step_number']}. {instruction['instruction']} \n"
    recipe['md'] = recipe_md
    recipe['timestamp'] = str(datetime.now())
    st.markdown(recipe_md)
    st.write("")