Anubhutineo commited on
Commit
46d8a05
Β·
verified Β·
1 Parent(s): 56dcc59

Add AI Storyteller application

Browse files
Files changed (1) hide show
  1. app.py +151 -0
app.py ADDED
@@ -0,0 +1,151 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import transformers
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ import gradio as gr
5
+ import random
6
+ import re
7
+ from typing import Dict, List
8
+ import warnings
9
+ import os
10
+
11
+ warnings.filterwarnings('ignore')
12
+
13
+ # Set environment variables for better performance
14
+ os.environ['TOKENIZERS_PARALLELISM'] = 'false'
15
+
16
+ class AIStoryteller:
17
+ """AI Storyteller optimized for Hugging Face Spaces"""
18
+
19
+ def __init__(self):
20
+ self.model = None
21
+ self.tokenizer = None
22
+ self.model_loaded = False
23
+ self.load_model()
24
+
25
+ def load_model(self):
26
+ """Load the AI model"""
27
+ try:
28
+ print("πŸ“₯ Loading DistilGPT-2 model...")
29
+ model_name = "distilgpt2"
30
+
31
+ self.tokenizer = AutoTokenizer.from_pretrained(model_name)
32
+ self.model = AutoModelForCausalLM.from_pretrained(
33
+ model_name,
34
+ torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
35
+ )
36
+
37
+ if self.tokenizer.pad_token is None:
38
+ self.tokenizer.pad_token = self.tokenizer.eos_token
39
+
40
+ self.model_loaded = True
41
+ print("βœ… Model loaded successfully!")
42
+ return True
43
+
44
+ except Exception as e:
45
+ print(f"❌ Model loading failed: {e}")
46
+ return False
47
+
48
+ def generate_story(self, prompt, max_length=100):
49
+ """Generate a story from a prompt"""
50
+ if not self.model_loaded:
51
+ return "❌ Model not loaded. Please try again."
52
+
53
+ try:
54
+ inputs = self.tokenizer.encode(prompt, return_tensors='pt', max_length=256, truncation=True)
55
+
56
+ with torch.no_grad():
57
+ outputs = self.model.generate(
58
+ inputs,
59
+ max_length=min(inputs.shape[1] + max_length, 256),
60
+ temperature=0.8,
61
+ do_sample=True,
62
+ top_p=0.9,
63
+ top_k=50,
64
+ pad_token_id=self.tokenizer.eos_token_id,
65
+ no_repeat_ngram_size=2,
66
+ repetition_penalty=1.1
67
+ )
68
+
69
+ full_story = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
70
+
71
+ if full_story.startswith(prompt):
72
+ generated_part = full_story[len(prompt):].strip()
73
+ return f"{prompt} {generated_part}"
74
+ else:
75
+ return full_story
76
+
77
+ except Exception as e:
78
+ return f"❌ Error generating story: {str(e)}"
79
+
80
+ # Initialize the storyteller
81
+ print("πŸš€ Initializing AI Storyteller...")
82
+ storyteller = AIStoryteller()
83
+
84
+ def generate_story_interface(prompt, genre, story_length):
85
+ """Main interface function"""
86
+ if not prompt or not prompt.strip():
87
+ return "❌ Please enter a story prompt!"
88
+
89
+ genre_starters = {
90
+ "Fantasy": "In a realm of magic and wonder,",
91
+ "Sci-Fi": "In the distant future among the stars,",
92
+ "Mystery": "On a foggy night filled with secrets,",
93
+ "Horror": "In the darkness where nightmares dwell,",
94
+ "Romance": "When two hearts found each other,",
95
+ "Adventure": "On a daring quest for glory,",
96
+ "Comedy": "In a world of laughter and mishaps,",
97
+ "Drama": "In a tale of human emotion,"
98
+ }
99
+
100
+ if genre in genre_starters:
101
+ full_prompt = f"{genre_starters[genre]} {prompt.strip()}"
102
+ else:
103
+ full_prompt = prompt.strip()
104
+
105
+ return storyteller.generate_story(full_prompt, max_length=int(story_length))
106
+
107
+ # Create Gradio interface
108
+ interface = gr.Interface(
109
+ fn=generate_story_interface,
110
+ inputs=[
111
+ gr.Textbox(
112
+ label="πŸ“ Story Prompt",
113
+ placeholder="Enter your story idea (e.g., 'a detective finds a mysterious letter')",
114
+ lines=3
115
+ ),
116
+ gr.Dropdown(
117
+ choices=["Fantasy", "Sci-Fi", "Mystery", "Horror", "Romance", "Adventure", "Comedy", "Drama"],
118
+ label="🎭 Genre",
119
+ value="Fantasy"
120
+ ),
121
+ gr.Slider(
122
+ minimum=30,
123
+ maximum=120,
124
+ value=80,
125
+ label="πŸ“ Story Length"
126
+ )
127
+ ],
128
+ outputs=gr.Textbox(
129
+ label="πŸ“š Generated Story",
130
+ lines=8
131
+ ),
132
+ title="🎭 AI Storyteller",
133
+ description="""
134
+ 🌟 **Create Amazing Stories with AI!** 🌟
135
+
136
+ Enter a creative prompt, choose your favorite genre, and let AI craft a unique story for you!
137
+ Perfect for writers, students, and anyone who loves creative storytelling.
138
+ """,
139
+ examples=[
140
+ ["a young wizard discovers a hidden library", "Fantasy", 100],
141
+ ["a detective receives a cryptic phone call", "Mystery", 80],
142
+ ["robots develop feelings", "Sci-Fi", 90],
143
+ ["two strangers meet in a coffee shop", "Romance", 70],
144
+ ["an explorer finds a secret cave", "Adventure", 85]
145
+ ],
146
+ theme=gr.themes.Soft()
147
+ )
148
+
149
+ # Launch the app
150
+ if __name__ == "__main__":
151
+ interface.launch()