CosmickVisions commited on
Commit
567d64c
·
verified ·
1 Parent(s): cc713fd

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +218 -0
app.py ADDED
@@ -0,0 +1,218 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import groq
3
+ import os
4
+ import tempfile
5
+ import uuid
6
+ from dotenv import load_dotenv
7
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
8
+ from langchain.vectorstores import FAISS
9
+ from langchain.embeddings import HuggingFaceEmbeddings
10
+ import fitz # PyMuPDF
11
+ import base64
12
+ from PIL import Image
13
+ import io
14
+
15
+ # Load environment variables
16
+ load_dotenv()
17
+ client = groq.Client(api_key=os.getenv("GROQ_LEGAL_API_KEY"))
18
+ embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
19
+
20
+ # Directory to store FAISS indexes
21
+ FAISS_INDEX_DIR = "faiss_indexes_tech"
22
+ if not os.path.exists(FAISS_INDEX_DIR):
23
+ os.makedirs(FAISS_INDEX_DIR)
24
+
25
+ # Dictionary to store user-specific vectorstores
26
+ user_vectorstores = {}
27
+
28
+ # Custom CSS for Tech theme
29
+ custom_css = """
30
+ :root {
31
+ --primary-color: #008080; /* Teal */
32
+ --secondary-color: #006666; /* Dark Teal */
33
+ --light-background: #E0FFFF; /* Light Cyan */
34
+ --dark-text: #333333;
35
+ --white: #FFFFFF;
36
+ --border-color: #E5E7EB;
37
+ }
38
+ body { background-color: var(--light-background); font-family: 'Inter', sans-serif; }
39
+ .container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; }
40
+ .header { background-color: var(--white); border-bottom: 2px solid var(--border-color); padding: 15px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
41
+ .header-title { color: var(--secondary-color); font-size: 1.8rem; font-weight: 700; text-align: center; }
42
+ .header-subtitle { color: var(--dark-text); font-size: 1rem; text-align: center; margin-top: 5px; }
43
+ .chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; min-height: 500px; }
44
+ .message-user { background-color: var(--primary-color) !important; color: var(--white) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
45
+ .message-bot { background-color: #F0F0F0 !important; color: var(--dark-text) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
46
+ .input-area { background-color: var(--white) !important; border-top: 1px solid var(--border-color) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
47
+ .input-box { border: 1px solid var(--border-color) !important; border-radius: 24px !important; padding: 12px 16px !important; box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }
48
+ .send-btn { background-color: var(--secondary-color) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; }
49
+ .clear-btn { background-color: #F0F0F0 !important; border: 1px solid var(--border-color) !important; border-radius: 24px !important; color: var(--dark-text) !important; padding: 8px 16px !important; font-weight: 500 !important; }
50
+ .pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; padding: 20px; }
51
+ .pdf-viewer-image { max-width: 100%; height: auto; border: 1px solid var(--border-color); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
52
+ .stats-box { background-color: #E0F0F0; padding: 10px; border-radius: 8px; margin-top: 10px; }
53
+ """
54
+
55
+ # Function to process PDF files (unchanged)
56
+ def process_pdf(pdf_file):
57
+ if pdf_file is None:
58
+ return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0}
59
+ try:
60
+ session_id = str(uuid.uuid4())
61
+ with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
62
+ temp_file.write(pdf_file)
63
+ pdf_path = temp_file.name
64
+
65
+ doc = fitz.open(pdf_path)
66
+ texts = [page.get_text() for page in doc]
67
+ page_images = []
68
+ for page in doc:
69
+ pix = page.get_pixmap()
70
+ img_bytes = pix.tobytes("png")
71
+ img_base64 = base64.b64encode(img_bytes).decode("utf-8")
72
+ page_images.append(img_base64)
73
+ total_pages = len(doc)
74
+ total_words = sum(len(text.split()) for text in texts)
75
+ doc.close()
76
+
77
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
78
+ chunks = text_splitter.create_documents(texts)
79
+ vectorstore = FAISS.from_documents(chunks, embeddings)
80
+ index_path = os.path.join(FAISS_INDEX_DIR, session_id)
81
+ vectorstore.save_local(index_path)
82
+ user_vectorstores[session_id] = vectorstore
83
+
84
+ os.unlink(pdf_path)
85
+ pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words}
86
+ return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state
87
+ except Exception as e:
88
+ if "pdf_path" in locals() and os.path.exists(pdf_path):
89
+ os.unlink(pdf_path)
90
+ return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0}
91
+
92
+ # Function to generate chatbot responses with Tech theme
93
+ def generate_response(message, session_id, model_name, history):
94
+ if not message:
95
+ return history
96
+ try:
97
+ context = ""
98
+ if session_id and session_id in user_vectorstores:
99
+ vectorstore = user_vectorstores[session_id]
100
+ docs = vectorstore.similarity_search(message, k=3)
101
+ if docs:
102
+ context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs)
103
+ system_prompt = "You are a technical assistant specializing in analyzing tech manuals, whitepapers, and documentation."
104
+ if context:
105
+ system_prompt += " Use the following context to answer the question if relevant: " + context
106
+ completion = client.chat.completions.create(
107
+ model=model_name,
108
+ messages=[
109
+ {"role": "system", "content": system_prompt},
110
+ {"role": "user", "content": message}
111
+ ],
112
+ temperature=0.7,
113
+ max_tokens=1024
114
+ )
115
+ response = completion.choices[0].message.content
116
+ history.append((message, response))
117
+ return history
118
+ except Exception as e:
119
+ history.append((message, f"Error generating response: {str(e)}"))
120
+ return history
121
+
122
+ # Functions to update PDF viewer (unchanged)
123
+ def update_pdf_viewer(pdf_state):
124
+ if not pdf_state["total_pages"]:
125
+ return 0, None, "No PDF uploaded yet"
126
+ try:
127
+ img_data = base64.b64decode(pdf_state["page_images"][0])
128
+ img = Image.open(io.BytesIO(img_data))
129
+ return pdf_state["total_pages"], img, f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}"
130
+ except Exception as e:
131
+ print(f"Error decoding image: {e}")
132
+ return 0, None, "Error displaying PDF"
133
+
134
+ def update_image(page_num, pdf_state):
135
+ if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]:
136
+ return None
137
+ try:
138
+ img_data = base64.b64decode(pdf_state["page_images"][page_num - 1])
139
+ img = Image.open(io.BytesIO(img_data))
140
+ return img
141
+ except Exception as e:
142
+ print(f"Error decoding image: {e}")
143
+ return None
144
+
145
+ # Gradio interface
146
+ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
147
+ current_session_id = gr.State(None)
148
+ pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0})
149
+ gr.HTML("""
150
+ <div class="header">
151
+ <div class="header-title">Tech-Vision</div>
152
+ <div class="header-subtitle">Analyze technical documents with Groq's LLM API.</div>
153
+ </div>
154
+ """)
155
+ with gr.Row(elem_classes="container"):
156
+ with gr.Column(scale=1, min_width=300):
157
+ pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary")
158
+ upload_button = gr.Button("Process PDF", variant="primary")
159
+ pdf_status = gr.Markdown("No PDF uploaded yet")
160
+ model_dropdown = gr.Dropdown(
161
+ choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
162
+ value="llama3-70b-8192",
163
+ label="Select Groq Model"
164
+ )
165
+ with gr.Column(scale=2, min_width=600):
166
+ with gr.Tabs():
167
+ with gr.TabItem("PDF Viewer"):
168
+ with gr.Column(elem_classes="pdf-viewer-container"):
169
+ page_slider = gr.Slider(minimum=1, maximum=1, step=1, label="Page Number", value=1)
170
+ pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image")
171
+ stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box")
172
+
173
+ with gr.Row(elem_classes="container"):
174
+ with gr.Column(scale=2, min_width=600):
175
+ chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
176
+ with gr.Row():
177
+ msg = gr.Textbox(show_label=False, placeholder="Ask about your technical document...", scale=5)
178
+ send_btn = gr.Button("Send", scale=1)
179
+ clear_btn = gr.Button("Clear Conversation")
180
+
181
+ # Event Handlers (unchanged)
182
+ upload_button.click(
183
+ process_pdf,
184
+ inputs=[pdf_file],
185
+ outputs=[current_session_id, pdf_status, pdf_state]
186
+ ).then(
187
+ update_pdf_viewer,
188
+ inputs=[pdf_state],
189
+ outputs=[page_slider, pdf_image, stats_display]
190
+ )
191
+
192
+ msg.submit(
193
+ generate_response,
194
+ inputs=[msg, current_session_id, model_dropdown, chatbot],
195
+ outputs=[chatbot]
196
+ ).then(lambda: "", None, [msg])
197
+
198
+ send_btn.click(
199
+ generate_response,
200
+ inputs=[msg, current_session_id, model_dropdown, chatbot],
201
+ outputs=[chatbot]
202
+ ).then(lambda: "", None, [msg])
203
+
204
+ clear_btn.click(
205
+ lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet"),
206
+ None,
207
+ [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display]
208
+ )
209
+
210
+ page_slider.change(
211
+ update_image,
212
+ inputs=[page_slider, pdf_state],
213
+ outputs=[pdf_image]
214
+ )
215
+
216
+ # Launch the app
217
+ if __name__ == "__main__":
218
+ demo.launch()