Update app.py

app.py

@@ -1,5 +1,4 @@
 from transformers import FlaxAutoModelForSeq2SeqLM, AutoTokenizer
-import gradio as gr
 
 MODEL_NAME_OR_PATH = "flax-community/t5-recipe-generation"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH, use_fast=True)
@@ -12,7 +11,8 @@ generation_kwargs = {
     "no_repeat_ngram_size": 3,
     "do_sample": True,
     "top_k": 60,
-    "top_p": 0.95
+    "top_p": 0.95,
+    "num_return_sequences": 1
 }
 
 special_tokens = tokenizer.all_special_tokens
@@ -37,15 +37,14 @@ def target_postprocessing(texts, special_tokens):
         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]
+def generate_recipe(items):
+    inputs = [prefix + items]
     inputs = tokenizer(
         inputs,
         max_length=256,
         padding="max_length",
         truncation=True,
-        return_tensors="jax"
+        return_tensors="pt"
     )
     input_ids = inputs.input_ids
     attention_mask = inputs.attention_mask
@@ -54,21 +53,12 @@ def generation_function(texts):
         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[0]  # Return the first generated recipe as a string
-
-
-iface = gr.Interface(
-    fn=generation_function,
-    inputs="text",
-    outputs="text",
-    title="Recipe Generation",
-    description="Generate a recipe based on an input text."
-)
+    generated = output_ids.squeeze().tolist()
+    generated_recipe = tokenizer.batch_decode(generated, skip_special_tokens=False)
+    generated_recipe = target_postprocessing(generated_recipe, special_tokens)
+    return generated_recipe[0]
 
-if __name__ == "__main__":
-    iface.launch()
+# Example usage
+input_items = "apple, cucumber"
+generated_recipe = generate_recipe(input_items)
+print(generated_recipe)