Update app.py

app.py

@@ -4,6 +4,9 @@ import torch
 
 st.title('Генерация текста GPT-моделью')
 st.subheader('Это приложение показывает разницу в генерации текста моделью rugpt3small, обученной на документах общей тематики и этой же моделью, дообученной на анекдотах')
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 # Загружаем токенайзер модели
 from transformers import GPT2Tokenizer
 tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
@@ -16,17 +19,17 @@ model_init = GPT2LMHeadModel.from_pretrained(
     output_attentions = False,
     output_hidden_states = False,
 )
+model_init.to(device);
+# # Это обученная модель, в нее загружаем веса
+# model = GPT2LMHeadModel.from_pretrained(
+#     'sberbank-ai/rugpt3small_based_on_gpt2',
+#     output_attentions = False,
+#     output_hidden_states = False,
+# )
 
-# Это обученная модель, в нее загружаем веса
-model = GPT2LMHeadModel.from_pretrained(
-    'sberbank-ai/rugpt3small_based_on_gpt2',
-    output_attentions = False,
-    output_hidden_states = False,
-)
-
-m = torch.load('model.pt')
-model.load_state_dict(m)
-
+# m = torch.load('model.pt')
+# model.load_state_dict(m)
+# model.to(device);
 
 str = st.text_input('Введите 1-4 слова начала текста, и подождите минутку', 'Мужик спрашивает у официанта')
 
@@ -34,7 +37,7 @@ str = st.text_input('Введите 1-4 слова начала текста, и
 # prompt – строка, которую примет на вход и продолжит модель
 
 # токенизируем строку
-prompt = tokenizer.encode(str, return_tensors='pt')
+prompt = tokenizer.encode(str, return_tensors='pt').to(device)
 
 # out будет содержать результаты генерации в виде списка
 out1 = model_init.generate(
@@ -56,7 +59,7 @@ out1 = model_init.generate(
     no_repeat_ngram_size=3,
     # сколько вернуть генераций
     num_return_sequences=3,
-    ).numpy() #).cpu().numpy()
+    ).cpu().numpy() #).numpy()
 
 st.write('\n------------------\n')
 st.subheader('Тексты на модели, обученной документами всех тематик:')
@@ -70,26 +73,26 @@ for out_ in out1:
     # print(tokenizer.decode(out_))
 
 
-# дообученная модель
-with torch.inference_mode():
-    # prompt = 'Мужик спрашивает официанта'
-    # prompt = tokenizer.encode(str, return_tensors='pt')
-    out2 = model.generate(
-        input_ids=prompt,
-        max_length=150,
-        num_beams=1,
-        do_sample=True,
-        temperature=1., 
-        top_k=5,
-        top_p=0.6,
-        no_repeat_ngram_size=2,
-        num_return_sequences=3,
-        ).numpy() #).cpu().numpy()
+# # дообученная модель
+# with torch.inference_mode():
+#     # prompt = 'Мужик спрашивает официанта'
+#     # prompt = tokenizer.encode(str, return_tensors='pt')
+#     out2 = model.generate(
+#         input_ids=prompt,
+#         max_length=150,
+#         num_beams=1,
+#         do_sample=True,
+#         temperature=1., 
+#         top_k=5,
+#         top_p=0.6,
+#         no_repeat_ngram_size=2,
+#         num_return_sequences=3,
+#         ).numpy() #).cpu().numpy()
     
-    st.subheader('Тексты на модели, обученной документами всех тематик и дообученной анекдотами:')
-    n = 0
-    for out_ in out2:
-        n += 1
-        st.write(tokenizer.decode(out_).rpartition('.')[0],'.')
-        # print(textwrap.fill(tokenizer.decode(out_), 100), end='\n------------------\n')
-        st.write('\n------------------\n')
+#     st.subheader('Тексты на модели, обученной документами всех тематик и дообученной анекдотами:')
+#     n = 0
+#     for out_ in out2:
+#         n += 1
+#         st.write(tokenizer.decode(out_).rpartition('.')[0],'.')
+#         # print(textwrap.fill(tokenizer.decode(out_), 100), end='\n------------------\n')
+#         st.write('\n------------------\n')