Update app.py

app.py

@@ -3,7 +3,9 @@ import torch
 import tokenizers
 import streamlit as st
 import re
+
 from PIL import Image
+from huggingface_hub import hf_hub_download
 
 
 @st.cache(hash_funcs={tokenizers.Tokenizer: lambda _: None, tokenizers.AddedToken: lambda _: None, re.Pattern: lambda _: None}, allow_output_mutation=True, suppress_st_warning=True)
@@ -37,22 +39,30 @@ def predict(text, model, tokenizer, n_beams=5, temperature=2.5, top_p=0.8, lengt
     return list(map(tokenizer.decode, out))[0]
 
 
-model, tokenizer = get_model('sberbank-ai/rugpt3medium_based_on_gpt2', 'korzh-medium_best_eval_loss.bin')
+medium_model, medium_tokenizer = get_model('sberbank-ai/rugpt3medium_based_on_gpt2', 'korzh-medium_best_eval_loss.bin')
+large_model, large_tokenizer = get_model('sberbank-ai/rugpt3large_based_on_gpt2', 'korzh-large_best_eval_loss.bin')
 
 # st.title("NeuroKorzh")
 
 image = Image.open('korzh.jpg')
 st.image(image, caption='NeuroKorzh')
 
+option = st.selectbox('Model to be used', ('medium', 'large'))
+
 st.markdown("\n")
 
-text = st.text_input(label='Starting point for text generation', value='Что делать, Макс?')
+text = st.text_area(label='Starting point for text generation', value='Что делать, Макс?', height=200)
 button = st.button('Go')
 
 if button:
     #try:
     with st.spinner("Generation in progress"):
-        result = predict(text, model, tokenizer)
+        if option == 'medium':
+            result = predict(text, medium_model, medium_tokenizer)
+        elif option == 'large':
+            result = predict(text, large_model, large_tokenizer)
+        else:
+            raise st.error('Error in selectbox')
     
     #st.subheader('Max Korzh:')
     #lines = result.split('\n')