updated app

app.py

@@ -2,24 +2,33 @@ import streamlit as st
 from transformers import AutoTokenizer,AutoModelForSeq2SeqLM
 
 @st.cache(persist=True)
-def load_model(input_complex_sentence,model):
-	
-	base_path = "flax-community/"
-	model_path = base_path + model
-	print(model_path)
-	tokenizer = AutoTokenizer.from_pretrained(model_path)
-	model     = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+def load_model(input_complex_sentence,model, tokenizer):
 	
 	tokenized_sentence = tokenizer(input_complex_sentence,return_tensors="pt")
 	result = model.generate(tokenized_sentence['input_ids'],attention_mask = tokenized_sentence['attention_mask'],max_length=256,num_beams=5)
-	print(result)
 	generated_sentence = tokenizer.decode(result[0],skip_special_tokens=True)
 	
 	return generated_sentence
 
 def main():
 
-	st.title("Sentence Split in English using T5 variants")
+	t5_base_path = "flax-community/t5-base-wikisplit"
+	t5_base_tokenizer = AutoTokenizer.from_pretrained(t5_base_path)
+	t5_base_model     = AutoModelForSeq2SeqLM.from_pretrained(t5_base_path)
+
+	t5_v1_1_base_path = "flax-community/t5-v1_1-base-wikisplit"
+	t5_v1_1_base_tokenizer = AutoTokenizer.from_pretrained(t5_v1_1_base_path)
+	t5_v1_1_base_model     = AutoModelForSeq2SeqLM.from_pretrained(t5_v1_1_base_path)
+
+	byt5_base_path = "flax-community/byt5-base-wikisplit"
+	byt5_base_tokenizer = AutoTokenizer.from_pretrained(byt5_base_path)
+	byt5_base_model     = AutoModelForSeq2SeqLM.from_pretrained(byt5_base_path)
+
+	t5_large_path = "flax-community/t5-large-wikisplit"
+	t5_large_tokenizer = AutoTokenizer.from_pretrained(t5_large_path)
+	t5_large_model     = AutoModelForSeq2SeqLM.from_pretrained(t5_large_path)
+
+	st.title("✂️ Sentence Split in English using T5 variants")
 	st.write("Sentence Split is the task of dividing a long Sentence into multiple Sentences")
 	
 	model = st.sidebar.selectbox(
@@ -31,8 +40,15 @@ def main():
 	input_complex_sentence = st.text_area("Please type a long Sentence to split",example)
 
 	if st.button('Simplify'):
- 
-		generated_sentence = load_model(input_complex_sentence, model)
+		
+		if model=="t5-base-wikisplit":
+			generated_sentence = load_model(input_complex_sentence, t5_base_model, t5_base_tokenizer)
+		elif model=="t5-v1_1-base-wikisplit":
+			generated_sentence = load_model(input_complex_sentence, t5_v1_1_base_model, t5_v1_1_base_tokenizer)
+		elif model=="byt5-base-wikisplit":
+			generated_sentence = load_model(input_complex_sentence, byt5_base_model, byt5_base_tokenizer)
+		else:
+			generated_sentence = load_model(input_complex_sentence, t5_large_model, t5_large_tokenizer)
 		st.write(generated_sentence)