Spencer525 commited on
Commit
1828bc5
1 Parent(s): 4ffe616

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +71 -0
app.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import asyncio
4
+ from langchain_core.prompts import PromptTemplate
5
+ from langchain_community.output_parsers.rail_parser import GuardrailsOutputParser
6
+ from langchain_community.document_loaders import PyPDFLoader
7
+ from langchain_google_genai import ChatGoogleGenerativeAI
8
+ import google.generativeai as genai
9
+ from langchain.chains.question_answering import load_qa_chain
10
+ import torch
11
+ from transformers import AutoTokenizer, AutoModelForCausalLM
12
+
13
+ # Gemini initialization and PDF QA function
14
+ async def initialize_gemini(file_path, question):
15
+ genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
16
+ model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
17
+ prompt_template = """Answer the question as precise as possible using the provided context. If the answer is
18
+ not contained in the context, say "answer not available in context" \n\n
19
+ Context: \n {context}?\n
20
+ Question: \n {question} \n
21
+ Answer:
22
+ """
23
+ prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
24
+ if os.path.exists(file_path):
25
+ pdf_loader = PyPDFLoader(file_path)
26
+ pages = pdf_loader.load_and_split()
27
+ context = "\n".join(str(page.page_content) for page in pages[:100])
28
+ stuff_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
29
+ stuff_answer = stuff_chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
30
+ return stuff_answer['output_text']
31
+ else:
32
+ return "Error: Unable to process the document. Please ensure the PDF file is valid."
33
+
34
+ # Mistral model initialization
35
+ def initialize_mistral():
36
+ model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
37
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
38
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
39
+ dtype = torch.bfloat16
40
+ model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
41
+ return tokenizer, model
42
+
43
+ # Mistral text generation function
44
+ def generate_mistral_text(prompt, tokenizer, model):
45
+ inputs = tokenizer.encode(prompt, return_tensors='pt').to(model.device)
46
+ outputs = model.generate(inputs, max_length=100)
47
+ return tokenizer.decode(outputs[0])
48
+
49
+ # Initialize Mistral model
50
+ mistral_tokenizer, mistral_model = initialize_mistral()
51
+
52
+ # Gradio interface function
53
+ async def pdf_qa(file, question):
54
+ gemini_answer = await initialize_gemini(file.name, question)
55
+ mistral_prompt = f"Based on this answer: '{gemini_answer}', provide a brief summary:"
56
+ mistral_summary = generate_mistral_text(mistral_prompt, mistral_tokenizer, mistral_model)
57
+ return f"Gemini Answer:\n{gemini_answer}\n\nMistral Summary:\n{mistral_summary}"
58
+
59
+ # Define Gradio Interface
60
+ input_file = gr.File(label="Upload PDF File")
61
+ input_question = gr.Textbox(label="Ask about the document")
62
+ output_text = gr.Textbox(label="Answer and Summary")
63
+
64
+ # Create Gradio Interface
65
+ gr.Interface(
66
+ fn=pdf_qa,
67
+ inputs=[input_file, input_question],
68
+ outputs=output_text,
69
+ title="PDF Question Answering System with Gemini and Mistral",
70
+ description="Upload a PDF file, ask questions about the content, and get answers from Gemini with a summary from Mistral."
71
+ ).launch()