PierreBrunelle commited on
Commit
04e9ddb
·
verified ·
1 Parent(s): fc3ce66

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +168 -0
app.py ADDED
@@ -0,0 +1,168 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import pixeltable as pxt
3
+ import numpy as np
4
+ from datetime import datetime
5
+ from pixeltable.functions.huggingface import sentence_transformer
6
+ from pixeltable.functions import openai
7
+ import os
8
+
9
+ # Ensure OpenAI API key is set
10
+ if 'OPENAI_API_KEY' not in os.environ:
11
+ os.environ['OPENAI_API_KEY'] = input('Enter your OpenAI API key: ')
12
+
13
+ # Initialize Pixeltable
14
+ pxt.drop_dir('story_builder', force=True)
15
+ pxt.create_dir('story_builder')
16
+
17
+ # Create embedding function
18
+ @pxt.expr_udf
19
+ def embed_text(text: str) -> np.ndarray:
20
+ return sentence_transformer(text, model_id='all-MiniLM-L6-v2')
21
+
22
+ # Create a table to store story contributions
23
+ story_table = pxt.create_table(
24
+ 'story_builder.contributions',
25
+ {
26
+ 'contributor': pxt.StringType(),
27
+ 'content': pxt.StringType(),
28
+ 'timestamp': pxt.TimestampType(),
29
+ 'cumulative_story': pxt.StringType()
30
+ }
31
+ )
32
+
33
+ # Add an embedding index to the content column
34
+ story_table.add_embedding_index('content', string_embed=embed_text)
35
+
36
+ @pxt.udf
37
+ def generate_summary(story: str) -> list[dict]:
38
+ system_msg = "You are an expert summarizer. Provide a concise summary of the given story, highlighting key plot points and themes."
39
+ user_msg = f"Story: {story}\n\nSummarize this story:"
40
+ return [
41
+ {'role': 'system', 'content': system_msg},
42
+ {'role': 'user', 'content': user_msg}
43
+ ]
44
+
45
+ story_table['summary_prompt'] = generate_summary(story_table.cumulative_story)
46
+ story_table['summary_response'] = openai.chat_completions(
47
+ messages=story_table.summary_prompt,
48
+ model='gpt-3.5-turbo',
49
+ max_tokens=200
50
+ )
51
+
52
+ @pxt.udf
53
+ def generate_continuation(context: str) -> list[dict]:
54
+ system_msg = "You are a creative writer. Continue the story based on the given context. Write a paragraph that logically follow the provided content."
55
+ user_msg = f"Context: {context}\n\nContinue the story:"
56
+ return [
57
+ {'role': 'system', 'content': system_msg},
58
+ {'role': 'user', 'content': user_msg}
59
+ ]
60
+
61
+ story_table['continuation_prompt'] = generate_continuation(story_table.cumulative_story)
62
+ story_table['continuation_response'] = openai.chat_completions(
63
+ messages=story_table.continuation_prompt,
64
+ model='gpt-3.5-turbo',
65
+ max_tokens=50
66
+ )
67
+
68
+ # Function to get the current cumulative story
69
+ def get_current_story():
70
+ latest_entry = story_table.tail(1)
71
+ if len(latest_entry) > 0:
72
+ return latest_entry['cumulative_story'][0]
73
+ return ""
74
+
75
+ # Functions for Gradio interface
76
+ def add_contribution(contributor, content):
77
+ current_story = get_current_story()
78
+ new_cumulative_story = current_story + " " + content if current_story else content
79
+
80
+ story_table.insert([{
81
+ 'contributor': contributor,
82
+ 'content': content,
83
+ 'timestamp': datetime.now(),
84
+ 'cumulative_story': new_cumulative_story
85
+ }])
86
+ return "Contribution added successfully!", new_cumulative_story
87
+
88
+ def get_similar_parts(query, num_results=5):
89
+ sim = story_table.content.similarity(query)
90
+ results = story_table.order_by(sim, asc=False).limit(num_results).select(story_table.content, story_table.contributor, sim=sim).collect()
91
+ return results.to_pandas()
92
+
93
+ def generate_next_part():
94
+ continuation = story_table.select(continuation=story_table.continuation_response.choices[0].message.content).tail(1)['continuation'][0]
95
+ return continuation
96
+
97
+ def summarize_story():
98
+ summary = story_table.select(summary=story_table.summary_response.choices[0].message.content).tail(1)['summary'][0]
99
+ return summary
100
+
101
+ # Gradio interface
102
+ with gr.Blocks(theme=gr.themes.Base()) as demo:
103
+ gr.HTML(
104
+ """
105
+ <div style="text-align: left; margin-bottom: 1rem;">
106
+ <img src="https://raw.githubusercontent.com/pixeltable/pixeltable/main/docs/source/data/pixeltable-logo-large.png" alt="Pixeltable" style="max-width: 150px;" />
107
+ </div>
108
+ """
109
+ )
110
+
111
+ gr.Markdown(
112
+ """
113
+ # 📚 Collaborative Story Builder
114
+
115
+ Welcome to the Collaborative Story Builder! This app allows multiple users to contribute to a story,
116
+ building it incrementally. Pixeltable manages the data, enables similarity search, and helps generate
117
+ continuations and summaries.
118
+ """
119
+ )
120
+
121
+ with gr.Tabs():
122
+ with gr.TabItem("Contribute"):
123
+ with gr.Row():
124
+ with gr.Column(scale=2):
125
+ contributor = gr.Textbox(label="Your Name")
126
+ content = gr.Textbox(label="Your Contribution", lines=5)
127
+ submit_btn = gr.Button("Submit Contribution", variant="primary")
128
+ with gr.Column(scale=3):
129
+ status = gr.Textbox(label="Status")
130
+ current_story = gr.Textbox(label="Current Story", lines=10, interactive=False)
131
+
132
+ with gr.TabItem("Search & Generate"):
133
+ with gr.Row():
134
+ with gr.Column():
135
+ search_query = gr.Textbox(label="Search Current Contributions")
136
+ num_results = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Number of Results")
137
+ search_btn = gr.Button("Search", variant="secondary")
138
+ search_results = gr.Dataframe(
139
+ headers=["Content", "Contributor", "Similarity"],
140
+ label="Similar Parts"
141
+ )
142
+
143
+ with gr.Column():
144
+ generate_btn = gr.Button("Generate Next Part", variant="primary")
145
+ generated_part = gr.Textbox(label="Generated Continuation", lines=5)
146
+
147
+ with gr.TabItem("Summary"):
148
+ summarize_btn = gr.Button("Summarize Story", variant="primary")
149
+ summary = gr.Textbox(label="Story Summary", lines=8)
150
+
151
+ submit_btn.click(add_contribution, inputs=[contributor, content], outputs=[status, current_story])
152
+ search_btn.click(get_similar_parts, inputs=[search_query, num_results], outputs=[search_results])
153
+ generate_btn.click(generate_next_part, outputs=[generated_part])
154
+ summarize_btn.click(summarize_story, outputs=[summary])
155
+
156
+ gr.HTML(
157
+ """
158
+ <div style="text-align: center; margin-top: 1rem; padding-top: 1rem; border-top: 1px solid #ccc;">
159
+ <p style="margin: 0; color: #666; font-size: 0.8em;">
160
+ Powered by <a href="https://github.com/pixeltable/pixeltable" target="_blank" style="color: #F25022; text-decoration: none;">Pixeltable</a>
161
+ | <a href="https://github.com/pixeltable/pixeltable" target="_blank" style="color: #666; text-decoration: none;">GitHub</a>
162
+ </p>
163
+ </div>
164
+ """
165
+ )
166
+
167
+ if __name__ == "__main__":
168
+ demo.launch()