Update gpt.py

gpt.py

@@ -1,7 +1,7 @@
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import streamlit as st
 import torch
-
+import textwrap
 
 tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
 model = GPT2LMHeadModel.from_pretrained(
@@ -16,32 +16,52 @@ model.load_state_dict(torch.load('modelgpt.pt', map_location=torch.device('cpu')
 
 
 prompt = st.text_input('Введите текст prompt:')
-length = st.slider('Длина генерируемой последовательности:', 10, 1000, 15)
+length = st.slider('Длина генерируемой последовательности:', 8, 256, 15)
 num_samples = st.slider('Число генераций:', 1, 10, 1)
 temperature = st.slider('Температура:', 1.0, 10.0, 2.0)
+top_k = st.slider('Количество наиболее вероятных слов генерации:', 10, 200, 50)
+top_k = st.slider('Минимальная суммарная вероятность топовых слов:', 0.4, 1.0, 0.9)
 
 
+# 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)
 
-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
+#     return generated_texts
 
 
 if st.button('Сгенерировать текст'):
-    generated_texts = generate_text(model, tokenizer, prompt, length, num_samples, temperature)
-    for i, text in enumerate(generated_texts):
+    
+    with torch.inference_mode():
+    prompt = tokenizer.encode(prompt, return_tensors='pt').to(device)
+    out = model.generate(
+        input_ids=prompt,
+        max_length=length,
+        num_beams=num_samples,
+        do_sample=True,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        no_repeat_ngram_size=3,
+        num_return_sequences=3,
+        ).cpu().numpy()
+    for i, out_ in enumerate(out):
         st.write(f'Текст {i+1}:')
-        st.write(text)
+        st.write(textwrap.fill(tokenizer.decode(out_), 100))
+
+
+    
+    # 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)