update

pages/1_Demo_1.py

@@ -62,13 +62,13 @@ else:
                                                    do_sample=False, truncation=True)
 
             print(male_generation)
-            st.session_state['male_continuations'] = [gen['generated_text'].replace(prompt, '') for gen, prompt in
+            st.session_state['male_continuations'] = [gen[0]['generated_text'].replace(prompt, '') for gen, prompt in
                                                       zip(male_generation, st.session_state['male_prompts'])]
 
             st.write('Generating text for female prompts...')
             female_generation = GPT2.text_generation(st.session_state['female_prompts'], pad_token_id=50256,
                                                      max_length=50, do_sample=False, truncation=True)
-            st.session_state['female_continuations'] = [gen['generated_text'].replace(prompt, '') for gen, prompt in
+            st.session_state['female_continuations'] = [gen[0]['generated_text'].replace(prompt, '') for gen, prompt in
                                                         zip(female_generation, st.session_state['female_prompts'])]
 
             st.write('Generated {} male continuations'.format(len(st.session_state['male_continuations'])))