import gradio as gr import groq import os import tempfile import uuid from dotenv import load_dotenv from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.vectorstores import FAISS from langchain.embeddings import HuggingFaceEmbeddings import fitz # PyMuPDF import base64 from PIL import Image import io # Load environment variables load_dotenv() client = groq.Client(api_key=os.getenv("GROQ_LEGAL_API_KEY")) embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # Directory to store FAISS indexes FAISS_INDEX_DIR = "faiss_indexes_tech" if not os.path.exists(FAISS_INDEX_DIR): os.makedirs(FAISS_INDEX_DIR) # Dictionary to store user-specific vectorstores user_vectorstores = {} # Custom CSS for Tech theme custom_css = """ :root { --primary-color: #008080; /* Teal */ --secondary-color: #006666; /* Dark Teal */ --light-background: #E0FFFF; /* Light Cyan */ --dark-text: #333333; --white: #FFFFFF; --border-color: #E5E7EB; } body { background-color: var(--light-background); font-family: 'Inter', sans-serif; } .container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; } .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); } .header-title { color: var(--secondary-color); font-size: 1.8rem; font-weight: 700; text-align: center; } .header-subtitle { color: var(--dark-text); font-size: 1rem; text-align: center; margin-top: 5px; } .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; } .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; } .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; } .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; } .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; } .send-btn { background-color: var(--secondary-color) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; } .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; } .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; } .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); } .stats-box { background-color: #E0F0F0; padding: 10px; border-radius: 8px; margin-top: 10px; } """ # Function to process PDF files (unchanged) def process_pdf(pdf_file): if pdf_file is None: return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0} try: session_id = str(uuid.uuid4()) with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file: temp_file.write(pdf_file) pdf_path = temp_file.name doc = fitz.open(pdf_path) texts = [page.get_text() for page in doc] page_images = [] for page in doc: pix = page.get_pixmap() img_bytes = pix.tobytes("png") img_base64 = base64.b64encode(img_bytes).decode("utf-8") page_images.append(img_base64) total_pages = len(doc) total_words = sum(len(text.split()) for text in texts) doc.close() text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) chunks = text_splitter.create_documents(texts) vectorstore = FAISS.from_documents(chunks, embeddings) index_path = os.path.join(FAISS_INDEX_DIR, session_id) vectorstore.save_local(index_path) user_vectorstores[session_id] = vectorstore os.unlink(pdf_path) pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words} return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state except Exception as e: if "pdf_path" in locals() and os.path.exists(pdf_path): os.unlink(pdf_path) return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0} # Function to generate chatbot responses with Tech theme def generate_response(message, session_id, model_name, history): if not message: return history try: context = "" if session_id and session_id in user_vectorstores: vectorstore = user_vectorstores[session_id] docs = vectorstore.similarity_search(message, k=3) if docs: context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs) system_prompt = "You are a technical assistant specializing in analyzing tech manuals, whitepapers, and documentation." if context: system_prompt += " Use the following context to answer the question if relevant: " + context completion = client.chat.completions.create( model=model_name, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": message} ], temperature=0.7, max_tokens=1024 ) response = completion.choices[0].message.content history.append((message, response)) return history except Exception as e: history.append((message, f"Error generating response: {str(e)}")) return history # Functions to update PDF viewer (unchanged) def update_pdf_viewer(pdf_state): if not pdf_state["total_pages"]: return 0, None, "No PDF uploaded yet" try: img_data = base64.b64decode(pdf_state["page_images"][0]) img = Image.open(io.BytesIO(img_data)) return pdf_state["total_pages"], img, f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}" except Exception as e: print(f"Error decoding image: {e}") return 0, None, "Error displaying PDF" def update_image(page_num, pdf_state): if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]: return None try: img_data = base64.b64decode(pdf_state["page_images"][page_num - 1]) img = Image.open(io.BytesIO(img_data)) return img except Exception as e: print(f"Error decoding image: {e}") return None # Gradio interface with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo: current_session_id = gr.State(None) pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0}) gr.HTML("""
Tech-Vision
Analyze technical documents with Groq's LLM API.
""") with gr.Row(elem_classes="container"): with gr.Column(scale=1, min_width=300): pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary") upload_button = gr.Button("Process PDF", variant="primary") pdf_status = gr.Markdown("No PDF uploaded yet") model_dropdown = gr.Dropdown( choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"], value="llama3-70b-8192", label="Select Groq Model" ) with gr.Column(scale=2, min_width=600): with gr.Tabs(): with gr.TabItem("PDF Viewer"): with gr.Column(elem_classes="pdf-viewer-container"): page_slider = gr.Slider(minimum=1, maximum=1, step=1, label="Page Number", value=1) pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image") stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box") with gr.Row(elem_classes="container"): with gr.Column(scale=2, min_width=600): chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container") with gr.Row(): msg = gr.Textbox(show_label=False, placeholder="Ask about your technical document...", scale=5) send_btn = gr.Button("Send", scale=1) clear_btn = gr.Button("Clear Conversation") # Event Handlers (unchanged) upload_button.click( process_pdf, inputs=[pdf_file], outputs=[current_session_id, pdf_status, pdf_state] ).then( update_pdf_viewer, inputs=[pdf_state], outputs=[page_slider, pdf_image, stats_display] ) msg.submit( generate_response, inputs=[msg, current_session_id, model_dropdown, chatbot], outputs=[chatbot] ).then(lambda: "", None, [msg]) send_btn.click( generate_response, inputs=[msg, current_session_id, model_dropdown, chatbot], outputs=[chatbot] ).then(lambda: "", None, [msg]) clear_btn.click( lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet"), None, [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display] ) page_slider.change( update_image, inputs=[page_slider, pdf_state], outputs=[pdf_image] ) # Launch the app if __name__ == "__main__": demo.launch()