Update pages/gpt.py

pages/gpt.py

@@ -1,46 +1,48 @@
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
-import torch
 import streamlit as st
-
+import torch
+import textwrap
+import plotly.express as px
+​
+​
+​
+​
+tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
 model = GPT2LMHeadModel.from_pretrained(
     'sberbank-ai/rugpt3small_based_on_gpt2',
     output_attentions = False,
     output_hidden_states = False,
 )
-
-tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
-
-
 # Вешаем сохраненные веса на нашу модель
-model.load_state_dict(torch.load('models/model.pt', map_location=torch.device('cpu')))
-
-prompt = st.text_input('Введите текст prompt:')
-length = st.slider('Длина генерируемой последовательности:', 10, 1000, 50)
-num_samples = st.slider('Число генераций:', 1, 10, 1)
-temperature = st.slider('Температура:', 0.1, 1.0, 0.5)
-
-
-
-def generate_text(model, tokenizer, prompt, length, num_samples, temperature):
-    input_ids = tokenizer.encode(prompt, return_tensors='pt')
-    output_sequences = model.generate(
-        input_ids=input_ids,
-        max_length=length,
-        num_return_sequences=num_samples,
-        temperature=temperature
-    )
-
-    generated_texts = []
-    for output_sequence in output_sequences:
-        generated_text = tokenizer.decode(output_sequence, clean_up_tokenization_spaces=True)
-        generated_texts.append(generated_text)
-
-    return generated_texts
-
-
-if st.button('Сгенерировать текст'):
-    generated_texts = generate_text(model, tokenizer, prompt, length, num_samples, temperature)
-    for i, text in enumerate(generated_texts):
-        st.write(f'Текст {i+1}:')
-        st.write(text)
-
+model.load_state_dict(torch.load('models/modelgpt.pt', map_location=torch.device('cpu')))
+​
+​
+length = st.sidebar.slider('**Длина генерируемой последовательности:**', 8, 256, 15)
+num_samples = st.sidebar.slider('**Число генераций:**', 1, 10, 1)
+temperature = st.sidebar.slider('**Температура:**', 1.0, 10.0, 2.0)
+top_k = st.sidebar.slider('**Количество наиболее вероятных слов генерации:**', 10, 200, 50)
+top_p = st.sidebar.slider('**Минимальная суммарная вероятность топовых слов:**', 0.4, 1.0, 0.9)
+​
+​
+prompt = st.text_input('**Введите текст 👇:**')
+if st.button('**Сгенерировать текст**'):
+ 
+    with torch.inference_mode():
+        prompt = tokenizer.encode(prompt, return_tensors='pt')
+        out = model.generate(
+            input_ids=prompt,
+            max_length=length,
+            num_beams=8,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            no_repeat_ngram_size=3,
+            num_return_sequences=num_samples,
+            ).cpu().numpy()
+        st.write('**_Результат_** 👇')
+        for i, out_ in enumerate(out):
+​
+            with st.expander(f'Текст {i+1}:'):
+                st.write(textwrap.fill(tokenizer.decode(out_), 100))
+                st.image("pict/wow.png")