Preeti Dave commited on
Commit
17e8b8b
·
1 Parent(s): 8acadb6
Files changed (1) hide show
  1. app.py +20 -7
app.py CHANGED
@@ -1,18 +1,31 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
- # Load the pre-trained model for text generation (GPT-2)
5
- generator = pipeline("text-generation", model="gpt2")
6
 
7
- # Function to generate poems based on a prompt
8
- def generate_poem(prompt):
9
- poem = generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
 
 
 
 
 
 
 
 
 
 
10
  return poem.strip()
11
 
12
- # Gradio interface
13
  iface = gr.Interface(
14
  fn=generate_poem,
15
- inputs=gr.Textbox(lines=2, placeholder="Enter a prompt for your poem here...", label="Poem Prompt"),
 
 
 
16
  outputs="text",
17
  title="Poem Generator",
18
  description="Generate poems with AI. Provide a prompt and let the AI create a poem for you!",
 
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
+ # Load a poetry-specific model
5
+ generator = pipeline("text-generation", model="AlecHelbling/gpt2-poetry")
6
 
7
+ # Function to generate poems based on a prompt and form
8
+ def generate_poem(prompt, form):
9
+ # Add more specific prompt tuning based on the selected form
10
+ if form == "Sonnet":
11
+ prompt = f"Write a sonnet about {prompt}"
12
+ elif form == "Haiku":
13
+ prompt = f"Write a haiku about {prompt}"
14
+ elif form == "Limerick":
15
+ prompt = f"Write a limerick about {prompt}"
16
+ else: # Free Verse
17
+ prompt = f"Write a free verse poem about {prompt}"
18
+
19
+ poem = generator(prompt, max_length=150, num_return_sequences=1)['generated_text'] # Increased max_length
20
  return poem.strip()
21
 
22
+ # Gradio interface with dropdown for poetic forms
23
  iface = gr.Interface(
24
  fn=generate_poem,
25
+ inputs=[
26
+ gr.Textbox(lines=2, placeholder="Enter a prompt for your poem here...", label="Poem Prompt"),
27
+ gr.Dropdown(choices=["Sonnet", "Haiku", "Limerick", "Free Verse"], label="Poetic Form")
28
+ ],
29
  outputs="text",
30
  title="Poem Generator",
31
  description="Generate poems with AI. Provide a prompt and let the AI create a poem for you!",