Update app.py

app.py

@@ -1,25 +1,73 @@
-import streamlit as st
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+import gradio as gr
+from transformers import FlaxAutoModelForSeq2SeqLM, AutoTokenizer
 
 MODEL_NAME_OR_PATH = "flax-community/t5-recipe-generation"
-tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME_OR_PATH, use_fast=True)
-model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME_OR_PATH)
-
-def generate_recipe(input_items):
-    prefix = "items: "
-    input_text = prefix + input_items
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    output_ids = model.generate(input_ids)
-    generated_recipe = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return generated_recipe
-
-def main():
-    st.title("Recipe Generation")
-    input_items = st.text_area("Enter the recipe instructions:")
-    if st.button("Generate Recipe"):
-        generated_recipe = generate_recipe(input_items)
-        st.subheader("Generated Recipe:")
-        st.text(generated_recipe)
+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,
+    "do_sample": True,
+    "top_k": 60,
+    "top_p": 0.95,
+}
+
+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[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.",
+)
 
 if __name__ == "__main__":
-    main()
+    iface.launch()