Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import gradio as gr
 import asyncio
 from langchain_core.prompts import PromptTemplate
-from langchain_community.output_parsers.rail_parser import GuardrailsOutputParser
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_google_genai import ChatGoogleGenerativeAI
 import google.generativeai as genai
@@ -11,9 +10,8 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Gemini PDF QA System
-async def initialize(file_path, question):
+async def initialize_gemini(file_path, question):
     genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-    model = genai.GenerativeModel('gemini-pro')
     model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
     prompt_template = """Answer the question as precise as possible using the provided context. If the answer is
                           not contained in the context, say "answer not available in context" \n\n
@@ -27,49 +25,48 @@ async def initialize(file_path, question):
         pages = pdf_loader.load_and_split()
         context = "\n".join(str(page.page_content) for page in pages[:30])
         stuff_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
-        stuff_answer = await stuff_chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
+        stuff_answer = await stuff_chain.acall({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
         return stuff_answer['output_text']
     else:
         return "Error: Unable to process the document. Please ensure the PDF file is valid."
 
-async def pdf_qa(file, question):
-    answer = await initialize(file.name, question)
-    return answer
-
 # Mistral Text Completion
-def load_mistral_model():
-    model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    dtype = torch.bfloat16
-    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
-    return tokenizer, model
+class MistralModel:
+    def __init__(self):
+        self.model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.dtype = torch.bfloat16
+        self.model = AutoModelForCausalLM.from_pretrained(self.model_path, torch_dtype=self.dtype, device_map=self.device)
 
-def generate_text(prompt, max_length=50):
-    tokenizer, model = load_mistral_model()
-    inputs = tokenizer.encode(prompt, return_tensors='pt').to(model.device)
-    outputs = model.generate(inputs, max_length=max_length)
-    return tokenizer.decode(outputs[0])
+    def generate_text(self, prompt, max_length=50):
+        inputs = self.tokenizer.encode(prompt, return_tensors='pt').to(self.model.device)
+        outputs = self.model.generate(inputs, max_length=max_length)
+        return self.tokenizer.decode(outputs[0])
 
-# Gradio Interface
-def pdf_qa_wrapper(file, question):
-    return asyncio.run(pdf_qa(file, question))
+mistral_model = MistralModel()
 
+# Combined function for both models
+async def process_input(file, question):
+    gemini_answer = await initialize_gemini(file.name, question)
+    mistral_answer = mistral_model.generate_text(question)
+    return gemini_answer, mistral_answer
+
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Combined PDF QA and Text Completion System")
+    gr.Markdown("# PDF Question Answering and Text Completion System")
     
-    with gr.Tab("PDF Question Answering"):
-        input_file = gr.File(label="Upload PDF File")
-        input_question = gr.Textbox(label="Ask about the document")
-        output_text_gemini = gr.Textbox(label="Answer - GeminiPro")
-        pdf_qa_button = gr.Button("Ask Question")
+    input_file = gr.File(label="Upload PDF File")
+    input_question = gr.Textbox(label="Ask a question or provide a prompt")
+    process_button = gr.Button("Process")
     
-    with gr.Tab("Text Completion"):
-        input_prompt = gr.Textbox(label="Enter prompt for text completion")
-        output_text_mistral = gr.Textbox(label="Completed Text - Mistral")
-        complete_text_button = gr.Button("Complete Text")
+    output_text_gemini = gr.Textbox(label="Answer - Gemini")
+    output_text_mistral = gr.Textbox(label="Answer - Mistral")
 
-    pdf_qa_button.click(pdf_qa_wrapper, inputs=[input_file, input_question], outputs=output_text_gemini)
-    complete_text_button.click(generate_text, inputs=input_prompt, outputs=output_text_mistral)
+    process_button.click(
+        fn=process_input,
+        inputs=[input_file, input_question],
+        outputs=[output_text_gemini, output_text_mistral]
+    )
 
 demo.launch()