Update app.py

app.py

@@ -1,250 +1,87 @@
 import streamlit as st
 import pdfplumber
 import pytesseract
-import openai
-from openai import OpenAI
+from PIL import Image
 import json
 import pandas as pd
-import numpy as np
-from PIL import Image
 from io import BytesIO
 import time
-import traceback
-import os
-import hashlib
+from openai import OpenAI
 import groq
 
-class SyntheticDataGenerator:
-    def __init__(self):
-        self.SUPPORTED_MODELS = {
-            "Deepseek": {
-                "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key),
-                "models": ["deepseek-chat"],
-                "key_name": "DEEPSEEK_KEY"
-            },
-            "OpenAI": {
-                "client": lambda key: OpenAI(api_key=key),
-                "models": ["gpt-4-turbo"],
-                "key_name": "OPENAI_KEY"
-            },
-            "Mistral-Groq": {
-                "client": lambda key: groq.Groq(api_key=key),
-                "models": ["mixtral-8x7b-32768", "llama2-70b-4096"],
-                "key_name": "GROQ_KEY"
-            }
+class SyntheticDataFactory:
+    PROVIDER_CONFIG = {
+        "Deepseek": {
+            "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key),
+            "models": ["deepseek-chat"],
+            "key_label": "Deepseek API Key"
+        },
+        "OpenAI": {
+            "client": lambda key: OpenAI(api_key=key),
+            "models": ["gpt-4-turbo"],
+            "key_label": "OpenAI API Key"
+        },
+        "Groq": {
+            "client": lambda key: groq.Groq(api_key=key),
+            "models": ["mixtral-8x7b-32768", "llama2-70b-4096"],
+            "key_label": "Groq API Key"
         }
-        self.init_session()
+    }
 
-    def init_session(self):
-        if 'qa_pairs' not in st.session_state:
-            st.session_state.qa_pairs = []
-        if 'doc_data' not in st.session_state:
-            st.session_state.doc_data = []
+    def __init__(self):
+        self.init_session_state()
+
+    def init_session_state(self):
+        if 'qa_data' not in st.session_state:
+            st.session_state.qa_data = {
+                'pairs': [],
+                'metadata': {},
+                'exports': {}
+            }
         if 'processing' not in st.session_state:
             st.session_state.processing = {
                 'stage': 'idle',
-                'errors': [],
-                'warnings': []
+                'errors': []
             }
 
-    def process_pdf(self, uploaded_file):
-        """Robust PDF processing with advanced image handling"""
-        st.session_state.processing = {'stage': 'extracting', 'errors': [], 'warnings': []}
-        
-        try:
-            with pdfplumber.load(uploaded_file) as pdf:
-                for page_num, page in enumerate(pdf.pages, 1):
-                    page_data = self._process_page(page, page_num)
-                    st.session_state.doc_data.append(page_data)
-            
-            if len(st.session_state.processing['errors']) > 0:
-                st.error(f"Processed with {len(st.session_state.processing['errors'])} errors")
-            return True
-        except Exception as e:
-            self._log_error(f"PDF loading failed: {str(e)}")
-            return False
-
-    def _process_page(self, page, page_num):
-        """Process individual page with nested error handling"""
-        page_data = {"page": page_num, "text": "", "images": []}
-        
-        try:
-            page_data["text"] = page.extract_text() or ""
-        except Exception as e:
-            self._log_error(f"Page {page_num} text extraction failed: {str(e)}")
-
-        try:
-            for img_idx, img in enumerate(page.images):
-                img_data = self._process_image(img, page_num, img_idx)
-                if img_data:
-                    page_data["images"].append(img_data)
-        except Exception as e:
-            self._log_error(f"Page {page_num} image processing failed: {str(e)}")
-        
-        return page_data
-
-    def _process_image(self, img, page_num, img_idx):
-        """Advanced image processing with multiple fallbacks"""
-        try:
-            stream = img['stream']
-            width = self._get_dimension(stream, 'width')
-            height = self._get_dimension(stream, 'height')
-            
-            if width <= 0 or height <= 0:
-                raise ValueError("Invalid image dimensions")
-            
-            try:
-                return Image.frombytes("RGB", (width, height), stream.get_data())
-            except:
-                return Image.frombytes("L", (width, height), stream.get_data()).convert("RGB")
-        except Exception as e:
-            self._log_error(f"Page {page_num} image {img_idx} failed: {str(e)}")
-            return None
-
-    def _get_dimension(self, stream, dimension):
-        """Safe dimension extraction with multiple fallbacks"""
-        try:
-            return int(stream[dimension])
-        except:
-            try:
-                return int(stream['stream'][dimension])
-            except:
-                try:
-                    return int(stream['data'][dimension])
-                except:
-                    return 0
+    # Add remaining class methods from previous implementation
+    # (process_pdf, generate_qa, etc.)
 
-    def generate_qa(self, model_provider, model_name, temperature):
-        """Multi-model generation engine"""
-        st.session_state.processing = {'stage': 'generating', 'errors': []}
-        qa_pairs = []
-        
-        try:
-            client = self.SUPPORTED_MODELS[model_provider]["client"](
-                st.session_state[model_provider.lower() + "_key"]
-            )
-            
-            for page in st.session_state.doc_data:
-                content = self._get_page_content(page)
-                response = self._generate(client, model_name, content, temperature)
-                qa_pairs.extend(self._parse_response(response))
-            
-            st.session_state.qa_pairs = qa_pairs
-            return True
-        except Exception as e:
-            self._log_error(f"Generation failed: {str(e)}")
-            return False
-
-    def _generate(self, client, model, content, temp):
-        """Unified generation interface"""
-        if isinstance(client, groq.Groq):
-            return client.chat.completions.create(
-                messages=[{"role": "user", "content": content}],
-                model=model,
-                temperature=temp,
-                response_format={"type": "json_object"}
-            )
-        else:
-            return client.chat.completions.create(
-                model=model,
-                messages=[{"role": "user", "content": content}],
-                temperature=temp,
-                response_format={"type": "json_object"}
-            )
-
-    def _parse_response(self, response):
-        """Safe response parsing"""
-        try:
-            content = json.loads(response.choices[0].message.content)
-            return content.get('qa_pairs', [])
-        except Exception as e:
-            self._log_error(f"Response parsing failed: {str(e)}")
-            return []
-
-    def export_data(self, formats):
-        """Multi-format export system"""
-        exports = {}
-        df = pd.DataFrame(st.session_state.qa_pairs)
+def setup_sidebar():
+    """Configure sidebar with provider settings"""
+    with st.sidebar:
+        st.header("⚙️ AI Configuration")
+        provider = st.selectbox("Provider", list(SyntheticDataFactory.PROVIDER_CONFIG.keys()))
+        config = SyntheticDataFactory.PROVIDER_CONFIG[provider]
         
-        if 'JSON' in formats:
-            exports['synthetic_data.json'] = df.to_json(orient='records').encode()
-        if 'CSV' in formats:
-            exports['synthetic_data.csv'] = df.to_csv(index=False).encode()
-        if 'Parquet' in formats:
-            buffer = BytesIO()
-            df.to_parquet(buffer)
-            exports['synthetic_data.parquet'] = buffer.getvalue()
+        api_key = st.text_input(config["key_label"], type="password")
+        model = st.selectbox("Model", config["models"])
+        temp = st.slider("Temperature", 0.0, 1.0, 0.3)
         
-        return exports
-
-    def _log_error(self, message):
-        """Centralized error logging"""
-        st.session_state.processing['errors'].append(message)
-        st.error(message)
+        return provider, api_key, model, temp
 
-    def _get_page_content(self, page):
-        """Multimodal content extraction"""
-        text = page["text"]
-        if not text:
-            text = " ".join([pytesseract.image_to_string(img) for img in page["images"]])
-        return text
-
-def ui_setup():
-    """Enterprise-grade UI configuration"""
+def main():
     st.set_page_config(
-        page_title="Synthetic Data Factory Pro",
+        page_title="Enterprise Data Factory",
         page_icon="🏭",
-        layout="wide",
-        initial_sidebar_state="expanded"
+        layout="wide"
     )
     
-    with st.sidebar:
-        st.header("🔑 API Key Management")
-        for provider in ["Deepseek", "OpenAI", "Mistral-Groq"]:
-            st.text_input(
-                f"{provider} API Key",
-                type="password",
-                key=f"{provider.lower()}_key"
-            )
-        
-        st.header("🧠 AI Configuration")
-        provider = st.selectbox("Model Provider", ["Deepseek", "OpenAI", "Mistral-Groq"])
-        model = st.selectbox("Model", generator.SUPPORTED_MODELS[provider]["models"])
-        temp = st.slider("Temperature", 0.0, 1.0, 0.3)
+    # Initialize factory instance
+    factory = SyntheticDataFactory()
     
-    return provider, model, temp
-
-def main():
-    """Main application flow"""
-    provider, model, temp = ui_setup()
-    generator = SyntheticDataGenerator()
+    # Setup UI components
+    provider, api_key, model, temp = setup_sidebar()
     
-    st.title("🏭 Synthetic Data Factory Pro")
-    st.write("Enterprise-grade document processing with multi-modal AI")
+    st.title("🚀 Enterprise Synthetic Data Factory")
     
-    uploaded_file = st.file_uploader("Upload PDF Document", type=["pdf"])
+    # File upload and processing logic
+    uploaded_file = st.file_uploader("Upload Financial PDF", type=["pdf"])
     
-    if uploaded_file and st.button("Start Generation"):
-        if generator.process_pdf(uploaded_file):
-            if generator.generate_qa(provider, model, temp):
-                st.success("Generation completed successfully!")
-                
-                with st.expander("📊 Results Preview"):
-                    st.dataframe(pd.DataFrame(st.session_state.qa_pairs))
-                
-                with st.expander("📦 Advanced Export"):
-                    formats = st.multiselect(
-                        "Select formats",
-                        ["JSON", "CSV", "Parquet"],
-                        default=["JSON", "CSV"]
-                    )
-                    exports = generator.export_data(formats)
-                    
-                    if st.download_button("Export Package", 
-                                        data=json.dumps(exports),
-                                        file_name="synthetic_data.zip",
-                                        mime="application/zip"):
-                        st.success("Export package generated!")
+    if uploaded_file and api_key:
+        if st.button("Start Synthetic Generation"):
+            # Process document and generate Q&A pairs
+            pass  # Add processing logic here
 
 if __name__ == "__main__":
     main()