Update app.py

app.py

@@ -1,11 +1,9 @@
 import os
+import openai
 from pypdf import PdfReader
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_core.documents import Document
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_openai import ChatOpenAI, OpenAI
-from langchain_core.prompts import PromptTemplate
-from langchain.chains.summarize import load_summarize_chain
 import gradio as gr
 
 title = '''
@@ -15,7 +13,7 @@ title = '''
     How to Use:<br/>
     1. Upload a .PDF from your computer and fill OpenAI API key.<br/>
     2. Click the "Upload PDF" button, if successful a preview of your PDF text will be shown.<br/>
-    3. Click "Summarize!" and the output will be shown on the textbox bellow.<br/>
+    3. Click "Summarize!" and the output will be shown on the textbox below.<br/>
     You can also change some LLM configurations from the 'config' tab.<br/>
 </div>
 '''
@@ -23,31 +21,31 @@ title = '''
 desc_1 = '''
 <div style="text-align: left; font-family:Arial; color:Black; font-size: 14px;">
     <h3>Custom Prompt Template</h3>
-    <p style="text-align: left;">You can customize input prompt for the map and combine prompt of langchain's Map-Reduce Summarization pipeline 
-    using the texboxt bellow.<br/>
-    Prompt which will be fed into LLM use the format of : <b>{textbox input} + {pdf_text} +  "SUMMARY:"</b> <br/>
-    In essence each page of PDF will be summarized using map prompt, and each summary then be combined for final output using combine prompt.<br/>
+    <p style="text-align: left;">You can customize input prompt for the map and combine prompt of the map-reduce summarization pipeline 
+    using the textbox below.<br/>
+    Prompt which will be fed into LLM uses the format: <b>{textbox input} + {pdf_text} +  "SUMMARY:"</b> <br/>
+    In essence each page of PDF will be summarized using the map prompt, and each summary then be combined for final output using combine prompt.<br/>
     <a href="https://python.langchain.com/docs/use_cases/summarization">More Info on Map-Reduce for Summarization</a>
 </div>
 '''
 
 MAP_PROMPT = """
-                You will be given a page of text which section is enclosed in triple backticks (```).
-                Your goal is to give a summary of this section, ignoring references and footnote if present.
-                Your response should be at least 200 words only if input classified as academic text. 
-                Your response must fully encompass what was said in the page.
-
-                ```{text}```
-                SUMMARY:
-                """
+You will be given a page of text which section is enclosed in triple backticks (```).
+Your goal is to give a summary of this section, ignoring references and footnote if present.
+Your response should be at least 200 words only if input classified as academic text. 
+Your response must fully encompass what was said in the page.
+```{text}```
+SUMMARY:
+"""
+
 COMBINE_PROMPT = """
-                    Write a full summary of the following text enclosed in triple backticks (```).
-                    Full summary consists of a descriptive summary of at least 100 words (if possible),
-                    followed by numbered list which covers key points of the text.
+Write a full summary of the following text enclosed in triple backticks (```).
+Full summary consists of a descriptive summary of at least 100 words (if possible),
+followed by numbered list which covers key points of the text.
+```{text}```
+SUMMARY:
+"""
 
-                    ```{text}```
-                    SUMMARY:
-                    """
 config_info = {'temperature': 'Higher means more randomness to the output.',
                'max_tokens' : 'The maximum number of tokens to generate in the output.',
                'llm_list' : ''}
@@ -58,130 +56,150 @@ model_list = {'gpt-3.5-turbo':'chat',
 
 text_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n"], chunk_size=10000, chunk_overlap=250)
 
+# globals to hold parsed PDF
+pdf_docs = []
+page_count = 0
+
 def parse_pdf(pdf_file):
     global pdf_docs, page_count
+    if pdf_file is None:
+        raise gr.Error("Please upload a PDF file.")
     loader = PyPDFLoader(pdf_file.name)
     pdf_docs = loader.load_and_split(text_splitter)
     page_count = len(pdf_docs)
-    file_check(pdf_file)
-    return pdf_docs[0].page_content[:100]
-
-def summarize_pdf(api_key, model_name, temperature, llm_max_tokens, custom_map_prompt, custom_combine_prompt):
-    if not pdf_docs:
-        raise gr.Error("No PDF File Detected!")
-    
-    os.environ["OPENAI_API_KEY"] = api_key
-    
-    # Updated LLM Initialization
-    if model_list[model_name] == 'chat':
-        gpt_llm = ChatOpenAI(temperature=temperature, model=model_name, max_tokens=int(llm_max_tokens))
-    else:
-        gpt_llm = OpenAI(temperature=temperature, model=model_name, max_tokens=int(llm_max_tokens))
-
-    # Prompt Logic
-    map_template = PromptTemplate(template=generate_template(custom_map_prompt) if custom_map_prompt else MAP_PROMPT, input_variables=["text"])
-    combine_template = PromptTemplate(template=generate_template(custom_combine_prompt) if custom_combine_prompt else COMBINE_PROMPT, input_variables=["text"])
-
-    map_reduce_chain = load_summarize_chain(
-        gpt_llm,
-        chain_type="map_reduce",
-        map_prompt=map_template,
-        combine_prompt=combine_template,
-        token_max=3840 
-    )
-    
-    # Updated Invocation (invoke instead of __call__)
-    map_reduce_outputs = map_reduce_chain.invoke({"input_documents": pdf_docs})
-    return map_reduce_outputs['output_text']
+    file_check(pdf_file)   # will raise gr.Error on invalid
+    # show first 200 chars preview
+    return pdf_docs[0].page_content[:200] if page_count > 0 else ""
 
 def file_check(pdf_file):
-  if os.path.getsize(pdf_file.name)/1024 **2 > 1:
-    raise gr.Error("Maximum File Size is 1MB!")
-  elif page_count > 15:
-    raise gr.Error("Maximum File Length is 15 Pages!")
-  else:
-    pass
-      
-  # Build LLM Model
-  os.environ["OPENAI_API_KEY"] = api_key
-  if model_list[model_name] == 'chat':
-    gpt_llm = ChatOpenAI(temperature=temperature, model_name=model_name, max_tokens=int(llm_max_tokens))
-  else:
-    gpt_llm = OpenAI(temperature=temperature, model_name=model_name, max_tokens=int(llm_max_tokens))
-  
-  # Summarize PDF
-  if custom_map_prompt !="":
-    map_template = PromptTemplate(template=generate_template(custom_map_prompt), input_variables=["text"])
-  else:
-    map_template = PromptTemplate(template=MAP_PROMPT, input_variables=["text"])
-
-  if custom_combine_prompt !="":
-    combine_template = PromptTemplate(template=generate_template(custom_combine_prompt), input_variables=["text"])
-  else:
-    combine_template = PromptTemplate(template=COMBINE_PROMPT, input_variables=["text"])
-
-  map_reduce_chain = load_summarize_chain(
-    gpt_llm,
-    chain_type="map_reduce",
-    map_prompt=map_template,
-    combine_prompt=combine_template,
-    return_intermediate_steps=True,
-    token_max=3840 # limit the maximum number of tokens in the combined document (combine prompt).
-  )
-  map_reduce_outputs = map_reduce_chain({"input_documents": pdf_docs})
-  return map_reduce_outputs['output_text']
+    global page_count
+    size_mb = os.path.getsize(pdf_file.name) / (1024 ** 2)
+    if size_mb > 1:
+        raise gr.Error("Maximum File Size is 1MB!")
+    if page_count > 15:
+        raise gr.Error("Maximum File Length is 15 Pages!")
 
 def generate_template(custom_prompt):
-  custom_template = custom_prompt + '''
+    # ensure {text} placeholder remains for substitution
+    custom_template = custom_prompt + '''
+```{text}```
+SUMMARY:
+'''
+    return custom_template
+
+def _call_openai_chat(model, prompt, temperature, max_tokens):
+    # Chat model path
+    resp = openai.ChatCompletion.create(
+        model=model,
+        messages=[{"role": "user", "content": prompt}],
+        temperature=float(temperature),
+        max_tokens=int(max_tokens)
+    )
+    return resp["choices"][0]["message"]["content"].strip()
+
+def _call_openai_completion(model, prompt, temperature, max_tokens):
+    # Completion (instruct) model path
+    resp = openai.Completion.create(
+        model=model,
+        prompt=prompt,
+        temperature=float(temperature),
+        max_tokens=int(max_tokens),
+        n=1,
+        stop=None
+    )
+    return resp["choices"][0]["text"].strip()
 
-      ```{text}```
-      SUMMARY:
-      '''
-  return custom_template
+def summarize_pdf(api_key, model_name, temperature, llm_max_tokens, custom_map_prompt, custom_combine_prompt):
+    global pdf_docs
+    if not pdf_docs:
+        raise gr.Error("No PDF File Detected! Please upload a PDF first and click Upload.")
+
+    # set API key
+    openai.api_key = api_key
+    if not api_key:
+        raise gr.Error("OpenAI API key is required.")
+
+    model_type = model_list.get(model_name, "chat")
+    map_template = generate_template(custom_map_prompt) if custom_map_prompt else MAP_PROMPT
+    combine_template = generate_template(custom_combine_prompt) if custom_combine_prompt else COMBINE_PROMPT
+
+    map_summaries = []
+    try:
+        # MAP step: summarize each page/chunk
+        for idx, doc in enumerate(pdf_docs, start=1):
+            text = doc.page_content
+            prompt = map_template.replace("{text}", text)
+            if model_type == "chat":
+                summary = _call_openai_chat(model_name, prompt, temperature, llm_max_tokens)
+            else:
+                summary = _call_openai_completion(model_name, prompt, temperature, llm_max_tokens)
+            map_summaries.append(f"--- Page {idx} Summary ---\n{summary}\n")
+
+        # COMBINE step: combine map summaries and produce final summary
+        combined_text = "\n\n".join(map_summaries)
+        combine_prompt = combine_template.replace("{text}", combined_text)
+
+        # For combine, you might allow more tokens (we reuse llm_max_tokens here)
+        if model_type == "chat":
+            final = _call_openai_chat(model_name, combine_prompt, temperature, llm_max_tokens)
+        else:
+            final = _call_openai_completion(model_name, combine_prompt, temperature, llm_max_tokens)
+
+        return final
+    except openai.error.OpenAIError as e:
+        raise gr.Error(f"OpenAI API error: {str(e)}")
+    except Exception as e:
+        raise gr.Error(f"Unexpected error: {str(e)}")
 
 def main():
-  with gr.Blocks() as demo:
-    gr.HTML(title)
-    with gr.Tab("Main"):
-      with gr.Column():
-        pdf_doc = gr.File(label="Uploaded PDF:", file_types=['.pdf'])
-        with gr.Row():
-            submit_button = gr.Button(value="Upload!")
-            pdf_preview = gr.Textbox(label="PDF Preview:", lines=2, interactive=False)
-            
-        API_KEY = gr.Textbox(label="OpenAI API Key:", lines=1, type="password")
-        summarize_button = gr.Button(value="Summarize!")
-        summarized_text  = gr.Textbox(label="Summary", lines=10, show_copy_button=True)
-
-    with gr.Tab("Config"):
-      llm_model = gr.Dropdown(choices=model_list.keys(), label="LLM model used", value='gpt-3.5-turbo', interactive=True)
-      with gr.Row():
-        temperature = gr.Slider(minimum=0, maximum=0.5, step=0.1, label="temperature", info=config_info['temperature'])
-        llm_max_tokens = gr.Radio(choices=[128, 256, 512], value=256, interactive=True, label="LLM max tokens", info=config_info['max_tokens']) 
-      gr.HTML(desc_1)
-      with gr.Row():
-        user_map_prompt = gr.Textbox(label="Map PROMPT", lines=10, interactive=True)
-        user_comb_prompt = gr.Textbox(label="Combine PROMPT", lines=10, interactive=True)
-
-      with gr.Accordion("Default Template", open=False):
-        with gr.Row():
-          default_map_prompt = gr.Textbox(label="Map PROMPT", value=MAP_PROMPT, lines=10, interactive=False)
-          default_comb_prompt = gr.Textbox(label="Combine PROMPT", value=COMBINE_PROMPT, lines=10, interactive=False)
-      with gr.Accordion("User Custom Prompt Preview", open=False):
-        prompt_preview_button = gr.Button(value="View Custom Prompt")
-        with gr.Row():
-          custom_map_view = gr.Textbox(label="Map PROMPT", lines=10, interactive=False)
-          custom_comb_view = gr.Textbox(label="Combine PROMPT", lines=10, interactive=False)
-
-        prompt_preview_button.click(generate_template, inputs=[user_map_prompt], outputs=[custom_map_view])
-        prompt_preview_button.click(generate_template, inputs=[user_comb_prompt], outputs=[custom_comb_view])
-
-    inputs_list = [API_KEY, llm_model, temperature, llm_max_tokens, user_map_prompt, user_comb_prompt]
-
-    submit_button.click(parse_pdf, inputs=[pdf_doc], outputs=[pdf_preview])
-    summarize_button.click(summarize_pdf, inputs=inputs_list, outputs=[summarized_text])
-      
-  demo.queue(concurrency_count=1).launch(share=True)
+    with gr.Blocks() as demo:
+        gr.HTML(title)
+        with gr.Tab("Main"):
+            with gr.Column():
+                pdf_doc = gr.File(label="Uploaded PDF:", file_types=['.pdf'])
+                with gr.Row():
+                    submit_button = gr.Button(value="Upload!")
+                    pdf_preview = gr.Textbox(label="PDF Preview:", lines=4, interactive=False)
+
+                API_KEY = gr.Textbox(label="OpenAI API Key:", lines=1, type="password")
+                summarize_button = gr.Button(value="Summarize!")
+                summarized_text  = gr.Textbox(label="Summary", lines=10, show_copy_button=True)
+
+        with gr.Tab("Config"):
+            llm_model = gr.Dropdown(choices=list(model_list.keys()), label="LLM model used", value='gpt-3.5-turbo', interactive=True)
+            with gr.Row():
+                temperature = gr.Slider(minimum=0, maximum=0.5, step=0.1, label="temperature", info=config_info['temperature'])
+                llm_max_tokens = gr.Radio(choices=[128, 256, 512], value=256, interactive=True, label="LLM max tokens", info=config_info['max_tokens'])
+            gr.HTML(desc_1)
+            with gr.Row():
+                user_map_prompt = gr.Textbox(label="Map PROMPT", lines=10, interactive=True)
+                user_comb_prompt = gr.Textbox(label="Combine PROMPT", lines=10, interactive=True)
+
+            with gr.Accordion("Default Template", open=False):
+                with gr.Row():
+                    default_map_prompt = gr.Textbox(label="Map PROMPT", value=MAP_PROMPT, lines=10, interactive=False)
+                    default_comb_prompt = gr.Textbox(label="Combine PROMPT", value=COMBINE_PROMPT, lines=10, interactive=False)
+
+            with gr.Accordion("User Custom Prompt Preview", open=False):
+                prompt_preview_button = gr.Button(value="View Custom Prompt")
+                with gr.Row():
+                    custom_map_view = gr.Textbox(label="Map PROMPT", lines=10, interactive=False)
+                    custom_comb_view = gr.Textbox(label="Combine PROMPT", lines=10, interactive=False)
+
+                # preview custom templates (wrapping the user's prompt)
+                def preview_custom(p):
+                    if not p:
+                        return ""
+                    return generate_template(p)
+                prompt_preview_button.click(preview_custom, inputs=[user_map_prompt], outputs=[custom_map_view])
+                prompt_preview_button.click(preview_custom, inputs=[user_comb_prompt], outputs=[custom_comb_view])
+
+        inputs_list = [API_KEY, llm_model, temperature, llm_max_tokens, user_map_prompt, user_comb_prompt]
+
+        submit_button.click(parse_pdf, inputs=[pdf_doc], outputs=[pdf_preview])
+        summarize_button.click(summarize_pdf, inputs=inputs_list, outputs=[summarized_text])
+
+    demo.queue(concurrency_count=1).launch(share=True)
 
 if __name__ == "__main__":
     main()