Update app.py

app.py

@@ -6,17 +6,14 @@ import gradio as gr
 from PyPDF2 import PdfReader
 import logging
 import webbrowser
-from gradio_client import Client
+from huggingface_hub import InferenceClient
+from typing import Dict, List, Optional, Tuple
+import time
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Initialize Hugging Face models
-HUGGINGFACE_MODELS = {
-    "Phi-3 Mini 128k": "eswardivi/Phi-3-mini-128k-instruct",
-}
-
-# Common context window sizes
+# Constants
 CONTEXT_SIZES = {
     "4K": 4000,
     "8K": 8000,
@@ -25,15 +22,49 @@ CONTEXT_SIZES = {
     "200K": 200000
 }
 
-def copy_to_clipboard(text):
-    return text
+class ModelRegistry:
+    def __init__(self):
+        self.hf_models = {
+            "Phi-3 Mini 128k": "microsoft/Phi-3-mini-128k-instruct",
+            "Custom Model": ""
+        }
+        self.groq_models = self._fetch_groq_models()
 
-def open_chatgpt():
-    webbrowser.open('https://chat.openai.com/')
-    return "Opening ChatGPT in browser..."
+    def _fetch_groq_models(self) -> Dict[str, str]:
+        """Fetch available Groq models"""
+        try:
+            headers = {
+                "Authorization": f"Bearer {os.getenv('GROQ_API_KEY')}",
+                "Content-Type": "application/json"
+            }
+            response = requests.get("https://api.groq.com/openai/v1/models", headers=headers)
+            if response.status_code == 200:
+                models = response.json().get("data", [])
+                return {model["id"]: model["id"] for model in models}
+            else:
+                logging.error(f"Failed to fetch Groq models: {response.status_code}")
+                return self._get_default_groq_models()
+        except Exception as e:
+            logging.error(f"Error fetching Groq models: {e}")
+            return self._get_default_groq_models()
+
+    def _get_default_groq_models(self) -> Dict[str, str]:
+        """Return default Groq models when API is unavailable"""
+        return {
+            "llama-3.1-70b-versatile": "llama-3.1-70b-versatile",
+            "mixtral-8x7b-32768": "mixtral-8x7b-32768",
+            "llama-3.1-8b-instant": "llama-3.1-8b-instant"
+        }
+
+    def refresh_groq_models(self) -> Dict[str, str]:
+        """Refresh the list of available Groq models"""
+        self.groq_models = self._fetch_groq_models()
+        return self.groq_models
+
+# Initialize model registry
+model_registry = ModelRegistry()
 
-# Utility Functions
-def extract_text_from_pdf(pdf_path):
+def extract_text_from_pdf(pdf_path: str) -> str:
     """Extract text content from PDF file."""
     try:
         reader = PdfReader(pdf_path)
@@ -51,7 +82,7 @@ def extract_text_from_pdf(pdf_path):
         logging.error(f"Error reading PDF file: {e}")
         return f"Error reading PDF file: {e}"
 
-def format_content(text, format_type):
+def format_content(text: str, format_type: str) -> str:
     """Format extracted text according to specified format."""
     if format_type == 'txt':
         return text
@@ -65,7 +96,7 @@ def format_content(text, format_type):
         logging.error(f"Unsupported format: {format_type}")
         return f"Unsupported format: {format_type}"
 
-def split_into_snippets(text, context_size):
+def split_into_snippets(text: str, context_size: int) -> List[str]:
     """Split text into manageable snippets based on context size."""
     sentences = re.split(r'(?<=[.!?]) +', text)
     snippets = []
@@ -87,7 +118,7 @@ def split_into_snippets(text, context_size):
 
     return snippets
 
-def build_prompts(snippets, prompt_instruction, custom_prompt, snippet_num=None):
+def build_prompts(snippets: List[str], prompt_instruction: str, custom_prompt: Optional[str], snippet_num: Optional[int] = None) -> str:
     """Build formatted prompts from text snippets."""
     if snippet_num is not None:
         if 1 <= snippet_num <= len(snippets):
@@ -111,27 +142,120 @@ def build_prompts(snippets, prompt_instruction, custom_prompt, snippet_num=None)
     
     return "\n\n".join(prompts)
 
-def send_to_huggingface(prompt, model_name):
-    """Send prompt to Hugging Face model using gradio_client."""
+def send_to_hf_inference(prompt: str, model_name: str, api_key: str) -> str:
+    """Send prompt to HuggingFace using Inference API"""
     try:
-        client = Client(model_name)
-        response = client.predict(
-            prompt,  # message
-            0.9,    # temperature
-            True,   # sampling
-            512,    # max_new_tokens
-            api_name="/chat"
+        client = InferenceClient(api_key=api_key)
+        messages = [{"role": "user", "content": prompt}]
+        completion = client.chat.completions.create(
+            model=model_name,
+            messages=messages,
+            max_tokens=500
         )
-        return response
+        return completion.choices[0].message.content
+    except Exception as e:
+        logging.error(f"Error with HF inference: {e}")
+        return f"Error with HF inference: {e}"
+
+def send_to_groq(prompt: str, model_name: str, api_key: str) -> str:
+    """Send prompt to Groq API"""
+    try:
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json"
+        }
+        data = {
+            "model": model_name,
+            "messages": [{"role": "user", "content": prompt}]
+        }
+        response = requests.post(
+            "https://api.groq.com/openai/v1/chat/completions",
+            headers=headers,
+            json=data
+        )
+        return response.json()["choices"][0]["message"]["content"]
+    except Exception as e:
+        logging.error(f"Error with Groq API: {e}")
+        return f"Error with Groq API: {e}"
+
+def copy_to_clipboard(text: str) -> str:
+    """Copy text to clipboard"""
+    return "Text copied to clipboard!"
+
+def open_chatgpt() -> str:
+    """Open ChatGPT in browser"""
+    webbrowser.open('https://chat.openai.com/')
+    return "Opening ChatGPT in browser..."
+
+def process_pdf(pdf, fmt, ctx_size, snippet_num, prompt, model_selection, 
+                hf_model_choice, hf_custom_model, hf_api_key,
+                groq_model_choice, groq_api_key) -> Tuple[str, str, str, List[str]]:
+    """Process PDF and generate summary"""
+    try:
+        if not pdf:
+            return "Please upload a PDF file.", "", "", []
+        
+        # Extract text
+        text = extract_text_from_pdf(pdf.name)
+        if text.startswith("Error"):
+            return text, "", "", []
+        
+        # Format content
+        formatted_text = format_content(text, fmt)
+        
+        # Split into snippets
+        snippets = split_into_snippets(formatted_text, ctx_size)
+        
+        # Build prompts
+        default_prompt = "Summarize the following text:"
+        full_prompt = build_prompts(snippets, default_prompt, prompt, snippet_num)
+        
+        if isinstance(full_prompt, str) and full_prompt.startswith("Error"):
+            return full_prompt, "", "", []
+        
+        # Process with selected model
+        if model_selection == "HuggingFace Inference":
+            if not hf_api_key:
+                return "HuggingFace API key required.", full_prompt, "", []
+            
+            model_id = hf_custom_model if hf_model_choice == "Custom Model" else model_registry.hf_models[hf_model_choice]
+            summary = send_to_hf_inference(full_prompt, model_id, hf_api_key)
+            
+        elif model_selection == "Groq API":
+            if not groq_api_key:
+                return "Groq API key required.", full_prompt, "", []
+                
+            summary = send_to_groq(full_prompt, groq_model_choice, groq_api_key)
+            
+        else:  # OpenAI ChatGPT
+            summary = "Please use the Copy Prompt button and paste into ChatGPT."
+        
+        # Save files for download
+        files_to_download = []
+        
+        with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as prompt_file:
+            prompt_file.write(full_prompt)
+            files_to_download.append(prompt_file.name)
+            
+        if summary != "Please use the Copy Prompt button and paste into ChatGPT.":
+            with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as summary_file:
+                summary_file.write(summary)
+                files_to_download.append(summary_file.name)
+        
+        return "Processing complete!", full_prompt, summary, files_to_download
+        
     except Exception as e:
-        logging.error(f"Error interacting with Hugging Face model: {e}")
-        return f"Error interacting with Hugging Face model: {e}"
+        logging.error(f"Error processing PDF: {e}")
+        return f"Error processing PDF: {str(e)}", "", "", []
 
 # Main Interface
 with gr.Blocks(theme=gr.themes.Default()) as demo:
+    # Store context size value
+    context_size_value = gr.State(value=32000)
+    
     # Header
     gr.Markdown("# 📄 Smart PDF Summarizer")
-    gr.Markdown("Upload a PDF document and get AI-powered summaries using OpenAI or Hugging Face models.")
+    gr.Markdown("Upload a PDF document and get AI-powered summaries using various AI models.")
     
     # Main Content
     with gr.Row():
@@ -151,9 +275,10 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             
             gr.Markdown("### Context Window Size")
             with gr.Row():
+                context_buttons = []
                 for size_name, size_value in CONTEXT_SIZES.items():
-                    if gr.Button(size_name).click:
-                        context_size.value = size_value
+                    btn = gr.Button(size_name)
+                    context_buttons.append((btn, size_value))
                         
             context_size = gr.Slider(
                 minimum=1000,
@@ -176,30 +301,43 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             )
             
             model_choice = gr.Radio(
-                choices=["OpenAI ChatGPT", "Hugging Face Model"],
+                choices=["OpenAI ChatGPT", "HuggingFace Inference", "Groq API"],
                 value="OpenAI ChatGPT",
                 label="🤖 Model Selection"
             )
             
-            hf_model = gr.Dropdown(
-                choices=list(HUGGINGFACE_MODELS.keys()),
-                label="🔧 Hugging Face Model",
-                visible=False
-            )
+            with gr.Column(visible=False) as hf_options:
+                hf_model = gr.Dropdown(
+                    choices=list(model_registry.hf_models.keys()),
+                    label="🔧 HuggingFace Model",
+                    value="Phi-3 Mini 128k"
+                )
+                hf_custom_model = gr.Textbox(
+                    label="Custom Model ID",
+                    placeholder="Enter custom model ID...",
+                    visible=False
+                )
+                hf_api_key = gr.Textbox(
+                    label="🔑 HuggingFace API Key",
+                    type="password"
+                )
             
-            # Authentication moved down
-            with gr.Row(visible=False) as auth_row:
-                openai_api_key = gr.Textbox(
-                    label="🔑 OpenAI API Key",
-                    type="password",
-                    placeholder="Enter your OpenAI API key (optional)"
+            with gr.Column(visible=False) as groq_options:
+                groq_model = gr.Dropdown(
+                    choices=list(model_registry.groq_models.keys()),
+                    label="🔧 Groq Model",
+                    value=list(model_registry.groq_models.keys())[0]
+                )
+                groq_refresh_btn = gr.Button("🔄 Refresh Models")
+                groq_api_key = gr.Textbox(
+                    label="🔑 Groq API Key",
+                    type="password"
                 )
 
         # Right Column - Output
         with gr.Column(scale=1):
-            with gr.Row():
-                process_button = gr.Button("🚀 Process PDF", variant="primary")
-                
+            process_button = gr.Button("🚀 Process PDF", variant="primary")
+            
             progress_status = gr.Textbox(
                 label="📊 Progress",
                 interactive=False
@@ -226,63 +364,51 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                 )
 
     # Event Handlers
-    def toggle_hf_model(choice):
-        return gr.update(visible=choice == "Hugging Face Model"), gr.update(visible=choice == "OpenAI ChatGPT")
+    def update_context_size(size):
+        return gr.update(value=size)
 
-    def process_pdf(pdf, fmt, ctx_size, snippet_num, prompt, model_selection, hf_model_choice):
-        try:
-            if not pdf:
-                return "Please upload a PDF file.", "", "", None
-            
-            # Extract text
-            text = extract_text_from_pdf(pdf.name)
-            if text.startswith("Error"):
-                return text, "", "", None
-            
-            # Format content
-            formatted_text = format_content(text, fmt)
-            
-            # Split into snippets
-            snippets = split_into_snippets(formatted_text, ctx_size)
-            
-            # Build prompts
-            default_prompt = "Summarize the following text:"
-            full_prompt = build_prompts(snippets, default_prompt, prompt, snippet_num)
-            
-            if isinstance(full_prompt, str) and full_prompt.startswith("Error"):
-                return full_prompt, "", "", None
-            
-            # Generate summary based on model choice
-            if model_selection == "Hugging Face Model":
-                summary = send_to_huggingface(full_prompt, HUGGINGFACE_MODELS[hf_model_choice])
-            else:
-                summary = "Please use the Copy Prompt button and paste into ChatGPT."
-            
-            # Save files for download
-            files_to_download = []
-            
-            with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as prompt_file:
-                prompt_file.write(full_prompt)
-                files_to_download.append(prompt_file.name)
-                
-            if summary != "Please use the Copy Prompt button and paste into ChatGPT.":
-                with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as summary_file:
-                    summary_file.write(summary)
-                    files_to_download.append(summary_file.name)
-            
-            return "Processing complete!", full_prompt, summary, files_to_download
-            
-        except Exception as e:
-            logging.error(f"Error processing PDF: {e}")
-            return f"Error processing PDF: {str(e)}", "", "", None
+    def toggle_model_options(choice):
+        return (
+            gr.update(visible=choice == "HuggingFace Inference"),
+            gr.update(visible=choice == "Groq API")
+        )
+
+def refresh_groq_models_list():
+        updated_models = model_registry.refresh_groq_models()
+        return gr.update(choices=list(updated_models.keys()))
+
+    def toggle_custom_model(model_name):
+        return gr.update(visible=model_name == "Custom Model")
 
     # Connect event handlers
     model_choice.change(
-        toggle_hf_model,
+        toggle_model_options,
         inputs=[model_choice],
-        outputs=[hf_model, auth_row]
+        outputs=[hf_options, groq_options]
     )
-    
+
+    for btn, size_value in context_buttons:
+        btn.click(
+            update_context_size,
+            inputs=[],
+            outputs=[context_size]
+        ).then(
+            lambda x=size_value: x,
+            None,
+            context_size
+        )
+
+    hf_model.change(
+        toggle_custom_model,
+        inputs=[hf_model],
+        outputs=[hf_custom_model]
+    )
+
+    groq_refresh_btn.click(
+        refresh_groq_models_list,
+        outputs=[groq_model]
+    )
+
     process_button.click(
         process_pdf,
         inputs=[
@@ -292,7 +418,11 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             snippet_number,
             custom_prompt,
             model_choice,
-            hf_model
+            hf_model,
+            hf_custom_model,
+            hf_api_key,
+            groq_model,
+            groq_api_key
         ],
         outputs=[
             progress_status,
@@ -301,19 +431,19 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             download_files
         ]
     )
-    
+
     copy_prompt_button.click(
         copy_to_clipboard,
         inputs=[generated_prompt],
         outputs=[progress_status]
     )
-    
+
     copy_summary_button.click(
         copy_to_clipboard,
         inputs=[summary_output],
         outputs=[progress_status]
     )
-    
+
     open_chatgpt_button.click(
         open_chatgpt,
         outputs=[progress_status]
@@ -325,7 +455,10 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
     1. Upload a PDF document
     2. Choose output format and context window size
     3. Select snippet number (default: 1) or enter custom prompt
-    4. Select between OpenAI ChatGPT or Hugging Face model
+    4. Select your preferred model:
+       - OpenAI ChatGPT: Manual copy/paste workflow
+       - HuggingFace Inference: Direct API integration
+       - Groq API: High-performance inference
     5. Click 'Process PDF' to generate summary
     6. Use 'Copy Prompt' and 'Open ChatGPT' for manual processing
     7. Download generated files as needed
@@ -334,6 +467,7 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
     - Support for multiple PDF formats
     - Flexible text formatting options
     - Predefined context window sizes (4K to 200K)
+    - Multiple model integrations
     - Copy to clipboard functionality
     - Direct ChatGPT integration
     - Downloadable outputs