Update app.py

app.py

@@ -191,20 +191,23 @@ elif selected_model == "Генерация текста GPT-моделью по
     tokenizer = GPT2Tokenizer.from_pretrained(path)
     
     if st.button('Сделать гороскоп'):
-        start_time = time.time()
-        input_ids = tokenizer.encode(user_text_input, return_tensors="pt").to(device)
-        model.eval()
-        with torch.no_grad():
-            out = model.generate(input_ids, 
-                                do_sample=True,
-                                num_beams=2,
-                                temperature=1.1,
-                                top_p=0.9,
-                                max_length=50,
-                                )
-
-        generated_text = list(map(tokenizer.decode, out))[0]
-        end_time = time.time()
-        prediction_time = end_time - start_time
-
-        st.write(f'Ваше предсказание: {generated_text}')
+        with st.spinner('Генерация текста...'):
+            start_time = time.time()
+            input_ids = tokenizer.encode(user_text_input, return_tensors="pt").to(device)
+            model.eval()
+            with torch.no_grad():
+                out = model.generate(
+                    input_ids,
+                    do_sample=True,
+                    num_beams=2,
+                    temperature=1.1,
+                    top_p=0.9,
+                    max_length=50,
+                )
+    
+            generated_text = tokenizer.decode(out[0], skip_special_tokens=True)
+            end_time = time.time()
+            prediction_time = end_time - start_time
+    
+        st.success('Готово!')
+        st.write({generated_text})