Update app.py

app.py

@@ -1,8 +1,8 @@
-from transformers import FlaxAutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import T5Tokenizer, FlaxT5ForConditionalGeneration
 
-MODEL_NAME_OR_PATH = "flax-community/t5-recipe-generation"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH, use_fast=True)
-model = FlaxAutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME_OR_PATH)
+MODEL_NAME_OR_PATH = "t5-base"
+tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME_OR_PATH)
+model = FlaxT5ForConditionalGeneration.from_pretrained(MODEL_NAME_OR_PATH)
 
 prefix = "items: "
 generation_kwargs = {
@@ -15,12 +15,6 @@ generation_kwargs = {
     "num_return_sequences": 1
 }
 
-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, "")
@@ -32,8 +26,6 @@ def target_postprocessing(texts, special_tokens):
     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
 
@@ -44,18 +36,18 @@ def generate_recipe(items):
         max_length=256,
         padding="max_length",
         truncation=True,
-        return_tensors="pt"
+        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.squeeze().tolist()
-    generated_recipe = tokenizer.batch_decode(generated, skip_special_tokens=False)
-    generated_recipe = target_postprocessing(generated_recipe, special_tokens)
+    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]
 
 # Example usage