Update app.py

app.py

@@ -1,68 +1,25 @@
 import streamlit as st
-from transformers import T5Tokenizer, FlaxT5ForConditionalGeneration
-import numpy as np
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-MODEL_NAME_OR_PATH = "t5-base"
-tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME_OR_PATH)
-model = FlaxT5ForConditionalGeneration.from_pretrained(MODEL_NAME_OR_PATH)
+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)
 
-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,
-    "num_return_sequences": 1
-}
-
-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)
-        new_texts.append(text)
-    return new_texts
-
-def generate_recipe(items):
-    inputs = [prefix + items]
-    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
-    )
-    
-    # Convert output IDs to numpy array
-    output_ids = np.array(output_ids)
-
-    generated_recipe = tokenizer.batch_decode(output_ids, skip_special_tokens=False)
-    generated_recipe = target_postprocessing(generated_recipe, tokenizer.all_special_tokens)
-    return generated_recipe[0]
+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")
-
-    items = st.text_input("Enter food items separated by comma (e.g., apple, cucumber):")
+    input_items = st.text_area("Enter the recipe instructions:")
     if st.button("Generate Recipe"):
-        generated_recipe = generate_recipe(items)
-        st.write(generated_recipe)
+        generated_recipe = generate_recipe(input_items)
+        st.subheader("Generated Recipe:")
+        st.text(generated_recipe)
 
 if __name__ == "__main__":
     main()