t5 model get func

app.py

@@ -8,12 +8,19 @@ from PIL import Image
 
 
 @st.cache(hash_funcs={tokenizers.Tokenizer: lambda _: None, tokenizers.AddedToken: lambda _: None, re.Pattern: lambda _: None}, allow_output_mutation=True, suppress_st_warning=True)
-def get_model(model_name, model_path):
+def get_model_gpt(model_name):
     tokenizer = transformers.GPT2Tokenizer.from_pretrained(model_name)
     model = transformers.GPT2LMHeadModel.from_pretrained(model_name)
     model.eval()
     return model, tokenizer
 
+@st.cache(hash_funcs={tokenizers.Tokenizer: lambda _: None, tokenizers.AddedToken: lambda _: None, re.Pattern: lambda _: None}, allow_output_mutation=True, suppress_st_warning=True)
+def get_model_t5(model_name):
+    tokenizer = transformers.T5Tokenizer.from_pretrained(model_name)
+    model = transformers.T5ForConditionalGeneration.from_pretrained(model_name)
+    model.eval()
+    return model, tokenizer
+
 
 def predict_gpt(text, model, tokenizer, temperature):
     input_ids = tokenizer.encode(text, return_tensors="pt")
@@ -54,10 +61,8 @@ def predict_t5(text, model, tokenizer, temperature):
     generated_text = list(map(decode, out['sequences']))[0]
     return generated_text
 
-gpt_model, gpt_tokenizer = get_model('mipatov/rugpt3_nb_descr', 'mipatov/rugpt3_nb_descr')
-t5_model, t5_tokenizer = get_model('mipatov/rut5_nb_descr', 'mipatov/rut5_nb_descr')
-
-# st.title("NeuroKorzh")
+gpt_model, gpt_tokenizer = get_model_gpt('mipatov/rugpt3_nb_descr', 'mipatov/rugpt3_nb_descr')
+t5_model, t5_tokenizer = get_model_t5('mipatov/rut5_nb_descr', 'mipatov/rut5_nb_descr')
 
 
 option = st.selectbox('Выберите модель', ('GPT', 'T5'))