Update app.py

app.py

@@ -1,8 +1,6 @@
 import json
-
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
 from format.format_output import format_output
 from validate.validate_ingredients import validate_ingredients
 from device.get_device_id import get_device_id
@@ -11,34 +9,56 @@ tokenizer = AutoTokenizer.from_pretrained("Ashikan/dut-recipe-generator")
 model = AutoModelForCausalLM.from_pretrained("Ashikan/dut-recipe-generator")
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=get_device_id())
 
-def perform_model_inference(ingredients_list):
-    for ingredient_index in range(len(ingredients_list)):
-        ingredients_list[ingredient_index] = ingredients_list[ingredient_index].strip()
+def perform_model_inference(ingredients_list=None, recipe_name=None):
+    if ingredients_list:
+        for ingredient_index in range(len(ingredients_list)):
+            ingredients_list[ingredient_index] = ingredients_list[ingredient_index].strip()
 
-    input_text = '{"prompt": ' + json.dumps(ingredients_list)
+        input_text = '{"prompt": ' + json.dumps(ingredients_list)
+    elif recipe_name:
+        input_text = '{"prompt": "Generate ingredients and method for the recipe: ' + recipe_name + '"}'
+    else:
+        return "Invalid input"
 
     output = pipe(input_text, max_length=1024, temperature=0.1, do_sample=True, truncation=True)[0]["generated_text"]
 
     return format_output(output)
 
-def chat_function(history, user_input):
-    # Extract user ingredients
-    ingredients_list = user_input.lower().split(',')
+def chat_function(history, user_input, mode):
+    # If mode is "ingredients", process as ingredient list
+    if mode == "ingredients":
+        ingredients_list = user_input.lower().split(',')
+        
+        # Validate the ingredients
+        if len(ingredients_list) < 2:
+            error_text = "Please provide at least 2 ingredients, separated by commas."
+            history.append((user_input, error_text))
+            return history, ""
 
-    # Validate the ingredients
-    if len(ingredients_list) < 2:
-        error_text = "Please provide at least 2 ingredients, separated by commas."
-        history.append((user_input, error_text))
+        # Generate the recipe
+        history.append((user_input, "Generating recipe..."))
+        recipe = perform_model_inference(ingredients_list=ingredients_list)
+        history[-1] = (user_input, recipe)  # Replace the "Generating recipe..." message with the result
         return history, ""
 
-    # Generate the recipe
-    history.append((user_input, "Generating recipe..."))
-    recipe = perform_model_inference(ingredients_list)
-    history[-1] = (user_input, recipe)  # Replace the "Generating recipe..." message with the result
-    return history, ""
+    # If mode is "recipe", process as recipe name
+    elif mode == "recipe":
+        recipe_name = user_input.strip()
+
+        # Validate the recipe name
+        if not recipe_name:
+            error_text = "Please provide a valid recipe name."
+            history.append((user_input, error_text))
+            return history, ""
+
+        # Generate ingredients and method
+        history.append((user_input, "Generating ingredients and method..."))
+        recipe_details = perform_model_inference(recipe_name=recipe_name)
+        history[-1] = (user_input, recipe_details)  # Replace the "Generating ingredients..." message with the result
+        return history, ""
 
 # Define the Gradio interface
-with gr.Blocks(css="""
+with gr.Blocks(css=""" 
 #chatbot-container {
     background-color: #f7f7f8;
     border-radius: 10px;
@@ -75,13 +95,21 @@ with gr.Blocks(css="""
 }
 """) as recipegen:
     gr.Markdown("# Recipe Generator")
-    gr.Markdown("An AI model attempting to produce healthier, diabetic-friendly recipes. Start by entering ingredients.")
+    gr.Markdown("An AI model attempting to produce healthier, diabetic-friendly recipes. Start by entering ingredients or a recipe name.")
 
     chatbot = gr.Chatbot(elem_id="chatbot-container")
-    user_input = gr.Textbox(placeholder="Enter ingredients (e.g., chicken, onions, garlic)", label="Your Input")
+    user_input = gr.Textbox(placeholder="Enter ingredients or recipe name", label="Your Input")
+    
+    # Dropdown for selecting input mode (ingredients or recipe name)
+    mode_selector = gr.Dropdown(
+        choices=["ingredients", "recipe"], 
+        label="Select Input Mode", 
+        value="ingredients"
+    )
+
     submit_button = gr.Button("Generate")
 
     # Link the chatbot and input to the chat function
-    submit_button.click(chat_function, inputs=[chatbot, user_input], outputs=[chatbot, user_input])
+    submit_button.click(chat_function, inputs=[chatbot, user_input, mode_selector], outputs=[chatbot, user_input])
 
 recipegen.launch(share=True)