DrishtiSharma commited on
Commit
3f3b6a0
·
verified ·
1 Parent(s): 03cd4eb

Create rough.py

Browse files
Files changed (1) hide show
  1. rough.py +365 -0
rough.py ADDED
@@ -0,0 +1,365 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import os
3
+ from openai import OpenAI
4
+ import tempfile
5
+ from langchain.chains import ConversationalRetrievalChain
6
+ from langchain_openai import ChatOpenAI, OpenAIEmbeddings
7
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
8
+ from langchain_community.vectorstores import Chroma
9
+ from langchain_community.document_loaders import (
10
+ PyPDFLoader,
11
+ TextLoader,
12
+ CSVLoader
13
+ )
14
+ from datetime import datetime
15
+ from pydub import AudioSegment
16
+ import pytz
17
+
18
+ from langchain.chains import ConversationalRetrievalChain
19
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
20
+ from langchain_openai import ChatOpenAI, OpenAIEmbeddings
21
+ from langchain_community.vectorstores import Chroma
22
+ from langchain_community.document_loaders import PyPDFLoader, TextLoader, CSVLoader
23
+ import os
24
+ import tempfile
25
+ from datetime import datetime
26
+ import pytz
27
+
28
+
29
+ class DocumentRAG:
30
+ def __init__(self):
31
+ self.document_store = None
32
+ self.qa_chain = None
33
+ self.document_summary = ""
34
+ self.chat_history = []
35
+ self.last_processed_time = None
36
+ self.api_key = os.getenv("OPENAI_API_KEY") # Fetch the API key from environment variable
37
+ self.init_time = datetime.now(pytz.UTC)
38
+
39
+ if not self.api_key:
40
+ raise ValueError("API Key not found. Make sure to set the 'OPENAI_API_KEY' environment variable.")
41
+
42
+ # Persistent directory for Chroma to avoid tenant-related errors
43
+ self.chroma_persist_dir = "./chroma_storage"
44
+ os.makedirs(self.chroma_persist_dir, exist_ok=True)
45
+
46
+ def process_documents(self, uploaded_files):
47
+ """Process uploaded files by saving them temporarily and extracting content."""
48
+ if not self.api_key:
49
+ return "Please set the OpenAI API key in the environment variables."
50
+ if not uploaded_files:
51
+ return "Please upload documents first."
52
+
53
+ try:
54
+ documents = []
55
+ for uploaded_file in uploaded_files:
56
+ # Save uploaded file to a temporary location
57
+ temp_file_path = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]).name
58
+ with open(temp_file_path, "wb") as temp_file:
59
+ temp_file.write(uploaded_file.read())
60
+
61
+ # Determine the loader based on the file type
62
+ if temp_file_path.endswith('.pdf'):
63
+ loader = PyPDFLoader(temp_file_path)
64
+ elif temp_file_path.endswith('.txt'):
65
+ loader = TextLoader(temp_file_path)
66
+ elif temp_file_path.endswith('.csv'):
67
+ loader = CSVLoader(temp_file_path)
68
+ else:
69
+ return f"Unsupported file type: {uploaded_file.name}"
70
+
71
+ # Load the documents
72
+ try:
73
+ documents.extend(loader.load())
74
+ except Exception as e:
75
+ return f"Error loading {uploaded_file.name}: {str(e)}"
76
+
77
+ if not documents:
78
+ return "No valid documents were processed. Please check your files."
79
+
80
+ # Split text for better processing
81
+ text_splitter = RecursiveCharacterTextSplitter(
82
+ chunk_size=1000,
83
+ chunk_overlap=200,
84
+ length_function=len
85
+ )
86
+ documents = text_splitter.split_documents(documents)
87
+
88
+ # Combine text for summary
89
+ combined_text = " ".join([doc.page_content for doc in documents])
90
+ self.document_summary = self.generate_summary(combined_text)
91
+
92
+ # Create embeddings and initialize retrieval chain
93
+ embeddings = OpenAIEmbeddings(api_key=self.api_key)
94
+ self.document_store = Chroma.from_documents(
95
+ documents,
96
+ embeddings,
97
+ persist_directory=self.chroma_persist_dir # Persistent directory for Chroma
98
+ )
99
+
100
+ self.qa_chain = ConversationalRetrievalChain.from_llm(
101
+ ChatOpenAI(temperature=0, model_name='gpt-4', api_key=self.api_key),
102
+ self.document_store.as_retriever(search_kwargs={'k': 6}),
103
+ return_source_documents=True,
104
+ verbose=False
105
+ )
106
+
107
+ self.last_processed_time = datetime.now(pytz.UTC)
108
+ return "Documents processed successfully!"
109
+ except Exception as e:
110
+ return f"Error processing documents: {str(e)}"
111
+
112
+ def generate_summary(self, text):
113
+ """Generate a summary of the provided text."""
114
+ if not self.api_key:
115
+ return "API Key not set. Please set it in the environment variables."
116
+ try:
117
+ client = OpenAI(api_key=self.api_key)
118
+ response = client.chat.completions.create(
119
+ model="gpt-4",
120
+ messages=[
121
+ {"role": "system", "content": "Summarize the document content concisely and provide 3-5 key points for discussion."},
122
+ {"role": "user", "content": text[:4000]}
123
+ ],
124
+ temperature=0.3
125
+ )
126
+ return response.choices[0].message.content
127
+ except Exception as e:
128
+ return f"Error generating summary: {str(e)}"
129
+
130
+ def create_podcast(self):
131
+ """Generate a podcast script and audio based on the document summary."""
132
+ if not self.document_summary:
133
+ return "Please process documents before generating a podcast.", None
134
+
135
+ if not self.api_key:
136
+ return "Please set the OpenAI API key in the environment variables.", None
137
+
138
+ try:
139
+ client = OpenAI(api_key=self.api_key)
140
+
141
+ # Generate podcast script
142
+ script_response = client.chat.completions.create(
143
+ model="gpt-4",
144
+ messages=[
145
+ {"role": "system", "content": "You are a professional podcast producer. Create a natural dialogue based on the provided document summary."},
146
+ {"role": "user", "content": f"""Based on the following document summary, create a 1-2 minute podcast script:
147
+ 1. Clearly label the dialogue as 'Host 1:' and 'Host 2:'
148
+ 2. Keep the content engaging and insightful.
149
+ 3. Use conversational language suitable for a podcast.
150
+ 4. Ensure the script has a clear opening and closing.
151
+ Document Summary: {self.document_summary}"""}
152
+ ],
153
+ temperature=0.7
154
+ )
155
+
156
+ script = script_response.choices[0].message.content
157
+ if not script:
158
+ return "Error: Failed to generate podcast script.", None
159
+
160
+ # Convert script to audio
161
+ final_audio = AudioSegment.empty()
162
+ is_first_speaker = True
163
+
164
+ lines = [line.strip() for line in script.split("\n") if line.strip()]
165
+ for line in lines:
166
+ if ":" not in line:
167
+ continue
168
+
169
+ speaker, text = line.split(":", 1)
170
+ if not text.strip():
171
+ continue
172
+
173
+ try:
174
+ voice = "nova" if is_first_speaker else "onyx"
175
+ audio_response = client.audio.speech.create(
176
+ model="tts-1",
177
+ voice=voice,
178
+ input=text.strip()
179
+ )
180
+
181
+ temp_audio_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
182
+ audio_response.stream_to_file(temp_audio_file.name)
183
+
184
+ segment = AudioSegment.from_file(temp_audio_file.name)
185
+ final_audio += segment
186
+ final_audio += AudioSegment.silent(duration=300)
187
+
188
+ is_first_speaker = not is_first_speaker
189
+ except Exception as e:
190
+ print(f"Error generating audio for line: {text}")
191
+ print(f"Details: {e}")
192
+ continue
193
+
194
+ if len(final_audio) == 0:
195
+ return "Error: No audio could be generated.", None
196
+
197
+ output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
198
+ final_audio.export(output_file, format="mp3")
199
+ return script, output_file
200
+
201
+ except Exception as e:
202
+ return f"Error generating podcast: {str(e)}", None
203
+
204
+ def generate_summary(self, text):
205
+ """Generate a summary of the provided text."""
206
+ if not self.api_key:
207
+ return "API Key not set. Please set it in the environment variables."
208
+ try:
209
+ client = OpenAI(api_key=self.api_key)
210
+ response = client.chat.completions.create(
211
+ model="gpt-4",
212
+ messages=[
213
+ {"role": "system", "content": "Summarize the document content concisely and provide 3-5 key points for discussion."},
214
+ {"role": "user", "content": text[:4000]}
215
+ ],
216
+ temperature=0.3
217
+ )
218
+ return response.choices[0].message.content
219
+ except Exception as e:
220
+ return f"Error generating summary: {str(e)}"
221
+
222
+ def handle_query(self, question, history):
223
+ """Handle user queries."""
224
+ if not self.qa_chain:
225
+ return history + [("System", "Please process the documents first.")]
226
+ try:
227
+ preface = """
228
+ Instruction: Respond in English. Be professional and concise, keeping the response under 300 words.
229
+ If you cannot provide an answer, say: "I am not sure about this question. Please try asking something else."
230
+ """
231
+ query = f"{preface}\nQuery: {question}"
232
+
233
+ result = self.qa_chain({
234
+ "question": query,
235
+ "chat_history": [(q, a) for q, a in history]
236
+ })
237
+
238
+ if "answer" not in result:
239
+ return history + [("System", "Sorry, an error occurred.")]
240
+
241
+ history.append((question, result["answer"]))
242
+ return history
243
+ except Exception as e:
244
+ return history + [("System", f"Error: {str(e)}")]
245
+
246
+ # Initialize RAG system in session state
247
+ if "rag_system" not in st.session_state:
248
+ st.session_state.rag_system = DocumentRAG()
249
+
250
+ # Sidebar
251
+ with st.sidebar:
252
+ st.title("About")
253
+ st.markdown(
254
+ """
255
+ This app is inspired by the [RAG_HW HuggingFace Space](https://huggingface.co/spaces/wint543/RAG_HW).
256
+ It allows users to upload documents, generate summaries, ask questions, and create podcasts.
257
+ """
258
+ )
259
+ st.markdown("### Steps:")
260
+ st.markdown("1. Upload documents.")
261
+ st.markdown("2. Generate summaries.")
262
+ st.markdown("3. Ask questions.")
263
+ st.markdown("4. Create podcasts.")
264
+
265
+ # Streamlit UI
266
+ # Sidebar
267
+ #with st.sidebar:
268
+ #st.title("About")
269
+ #st.markdown(
270
+ #"""
271
+ #This app is inspired by the [RAG_HW HuggingFace Space](https://huggingface.co/spaces/wint543/RAG_HW).
272
+ #It allows users to:
273
+ #1. Upload and process documents
274
+ #2. Generate summaries
275
+ #3. Ask questions
276
+ #4. Create podcasts
277
+ #"""
278
+ #)
279
+
280
+ # Main App
281
+ st.title("Document Analyzer & Podcast Generator")
282
+
283
+ # Step 1: Upload and Process Documents
284
+ st.subheader("Step 1: Upload and Process Documents")
285
+ uploaded_files = st.file_uploader("Upload files (PDF, TXT, CSV)", accept_multiple_files=True)
286
+
287
+ if st.button("Process Documents"):
288
+ if uploaded_files:
289
+ # Process the uploaded files
290
+ result = st.session_state.rag_system.process_documents(uploaded_files)
291
+ if "successfully" in result:
292
+ st.success(result)
293
+ else:
294
+ st.error(result)
295
+ else:
296
+ st.warning("No files uploaded.")
297
+
298
+ # Step 2: Generate Summaries
299
+ st.subheader("Step 2: Generate Summaries")
300
+ st.write("Select Summary Language:")
301
+ summary_language_options = ["English", "Hindi", "Spanish", "French", "German", "Chinese", "Japanese"]
302
+ cols = st.columns(len(summary_language_options))
303
+ # Default selected option
304
+ summary_language = st.radio(
305
+ "",
306
+ summary_language_options,
307
+ horizontal=True,
308
+ key="summary_language"
309
+ )
310
+
311
+ if st.session_state.rag_system.document_summary:
312
+ st.text_area("Document Summary", st.session_state.rag_system.document_summary, height=200)
313
+ else:
314
+ st.info("Please process documents first to generate summaries.")
315
+ if st.button("Generate Summary"):
316
+ st.session_state.rag_system.document_summary = st.session_state.rag_system.generate_summary(
317
+ st.session_state.rag_system.document_summary,
318
+ summary_language
319
+ )
320
+
321
+ # Step 3: Ask Questions
322
+ st.subheader("Step 3: Ask Questions")
323
+ st.write("Select Q&A Language:")
324
+ qa_language_options = ["English", "Hindi", "Spanish", "French", "German", "Chinese", "Japanese"]
325
+ qa_language = st.radio(
326
+ "",
327
+ qa_language_options,
328
+ horizontal=True,
329
+ key="qa_language"
330
+ )
331
+
332
+ if st.session_state.rag_system.qa_chain:
333
+ history = []
334
+ user_question = st.text_input("Ask a question:")
335
+ if st.button("Submit Question"):
336
+ # Handle the user query
337
+ history = st.session_state.rag_system.handle_query(user_question, history, qa_language)
338
+ for question, answer in history:
339
+ st.chat_message("user").write(question)
340
+ st.chat_message("assistant").write(answer)
341
+ else:
342
+ st.info("Please process documents first to enable Q&A.")
343
+
344
+ # Step 4: Generate Podcast
345
+ st.subheader("Step 4: Generate Podcast")
346
+ st.write("Select Podcast Language:")
347
+ podcast_language_options = ["English", "Hindi", "Spanish", "French", "German", "Chinese", "Japanese"]
348
+ podcast_language = st.radio(
349
+ "",
350
+ podcast_language_options,
351
+ horizontal=True,
352
+ key="podcast_language"
353
+ )
354
+
355
+ if st.session_state.rag_system.document_summary:
356
+ if st.button("Generate Podcast"):
357
+ script, audio_path = st.session_state.rag_system.create_podcast(podcast_language)
358
+ if audio_path:
359
+ st.text_area("Generated Podcast Script", script, height=200)
360
+ st.audio(audio_path, format="audio/mp3")
361
+ st.success("Podcast generated successfully! You can listen to it above.")
362
+ else:
363
+ st.error(script)
364
+ else:
365
+ st.info("Please process documents and generate summaries before creating a podcast.")