Update app.py

app.py

@@ -1,15 +1,11 @@
 import streamlit as st
 import random
 from transformers import pipeline
+import pandas as pd
 
-# Load the GPT2 model
 generator = pipeline('text-generation', model='gpt2')
-#generator = pipeline('text-generation', model='bloom-560m')
-
-# Set the maximum length of the generated prompt
 max_length = 50
 
-# Define the prompts and solutions
 prompts = {
     "Difficulty sleeping": [
         "Try keeping a consistent sleep schedule and avoid caffeine before bedtime.",
@@ -34,28 +30,25 @@ prompts = {
     ]
 }
 
-# Define the function to generate the prompts and solutions
 def generate_prompt(prompt):
-    # Get a random solution for the prompt
     solution = random.choice(prompts[prompt])
-    # Generate the prompt text
     prompt_text = f"What can I do to {prompt.lower()}? "
-    # Generate the prompt output
     output = generator(prompt_text, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)
-    # Decode the prompt output
     output_text = output[0]['generated_text'][len(prompt_text):].strip()
-    # Return the generated prompt and solution
-    return output_text, solution
+    return prompt_text, output_text, solution
 
-# Set the app title
 st.title('ICL-LM Interface')
-# Get the user input
 option = st.selectbox('Select a problem:', list(prompts.keys()))
 
 if st.button('Generate Prompt and Solution'):
-    # Generate the prompt and solution
-    prompt, solution = generate_prompt(option)
-    # Display the prompt
-    st.write('Prompt:', prompt)
-    # Display the solution
-    st.write('Solution:', solution)
+    results = []
+    for _ in range(3):
+        prompt_text, prompt, solution = generate_prompt(option)
+        results.append([prompt_text, prompt, solution])
+        
+    with open('results.txt', 'w') as f:
+        for result in results:
+            f.write(f"{result[0]}\t{result[1]}\t{result[2]}\n")
+    
+    df = pd.read_csv('results.txt', sep='\t', header=None, names=['Input', 'Prompt', 'Solution'])
+    st.write(df)