Roberta2024 commited on
Commit
4def369
1 Parent(s): 959d70e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -32
app.py CHANGED
@@ -1,6 +1,5 @@
1
  import os
2
  import gradio as gr
3
- import asyncio
4
  from langchain_core.prompts import PromptTemplate
5
  from langchain_community.document_loaders import PyPDFLoader
6
  from langchain_google_genai import ChatGoogleGenerativeAI
@@ -19,44 +18,48 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
19
  dtype = torch.bfloat16
20
  mistral_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
21
 
22
- async def initialize(file_path, question):
23
- model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
24
- prompt_template = """Answer the question as precise as possible using the provided context. If the answer is
25
- not contained in the context, say "answer not available in context" \n\n
26
- Context: \n {context}?\n
27
- Question: \n {question} \n
28
- Answer:
29
- """
30
- prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
31
-
32
- if os.path.exists(file_path):
33
- pdf_loader = PyPDFLoader(file_path)
34
- pages = pdf_loader.load_and_split()
35
- context = "\n".join(str(page.page_content) for page in pages[:30])
36
- stuff_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
37
- stuff_answer = await stuff_chain.acall({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
38
- gemini_answer = stuff_answer['output_text']
39
 
40
- # Use Mistral model for additional text generation
41
- mistral_prompt = f"Based on this answer: {gemini_answer}\nGenerate a follow-up question:"
42
- mistral_inputs = mistral_tokenizer.encode(mistral_prompt, return_tensors='pt').to(device)
43
- with torch.no_grad():
44
- mistral_outputs = mistral_model.generate(mistral_inputs, max_length=50)
45
- mistral_output = mistral_tokenizer.decode(mistral_outputs[0], skip_special_tokens=True)
46
-
47
- combined_output = f"Gemini Answer: {gemini_answer}\n\nMistral Follow-up: {mistral_output}"
48
- return combined_output
49
- else:
50
- return "Error: Unable to process the document. Please ensure the PDF file is valid."
 
 
 
 
 
 
 
 
 
 
51
 
52
  # Define Gradio Interface
53
  input_file = gr.File(label="Upload PDF File")
54
  input_question = gr.Textbox(label="Ask about the document")
55
  output_text = gr.Textbox(label="Answer - Combined Gemini and Mistral")
56
 
57
- async def pdf_qa(file, question):
58
- answer = await initialize(file.name, question)
59
- return answer
 
60
 
61
  # Create Gradio Interface
62
  gr.Interface(
 
1
  import os
2
  import gradio as gr
 
3
  from langchain_core.prompts import PromptTemplate
4
  from langchain_community.document_loaders import PyPDFLoader
5
  from langchain_google_genai import ChatGoogleGenerativeAI
 
18
  dtype = torch.bfloat16
19
  mistral_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
20
 
21
+ def initialize(file_path, question):
22
+ try:
23
+ model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
24
+ prompt_template = """Answer the question as precise as possible using the provided context. If the answer is
25
+ not contained in the context, say "answer not available in context" \n\n
26
+ Context: \n {context}?\n
27
+ Question: \n {question} \n
28
+ Answer:
29
+ """
30
+ prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
 
 
 
 
 
 
 
31
 
32
+ if os.path.exists(file_path):
33
+ pdf_loader = PyPDFLoader(file_path)
34
+ pages = pdf_loader.load_and_split()
35
+ context = "\n".join(str(page.page_content) for page in pages[:30])
36
+ stuff_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
37
+ stuff_answer = stuff_chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
38
+ gemini_answer = stuff_answer['output_text']
39
+
40
+ # Use Mistral model for additional text generation
41
+ mistral_prompt = f"Based on this answer: {gemini_answer}\nGenerate a follow-up question:"
42
+ mistral_inputs = mistral_tokenizer.encode(mistral_prompt, return_tensors='pt').to(device)
43
+ with torch.no_grad():
44
+ mistral_outputs = mistral_model.generate(mistral_inputs, max_length=50)
45
+ mistral_output = mistral_tokenizer.decode(mistral_outputs[0], skip_special_tokens=True)
46
+
47
+ combined_output = f"Gemini Answer: {gemini_answer}\n\nMistral Follow-up: {mistral_output}"
48
+ return combined_output
49
+ else:
50
+ return "Error: Unable to process the document. Please ensure the PDF file is valid."
51
+ except Exception as e:
52
+ return f"An error occurred: {str(e)}"
53
 
54
  # Define Gradio Interface
55
  input_file = gr.File(label="Upload PDF File")
56
  input_question = gr.Textbox(label="Ask about the document")
57
  output_text = gr.Textbox(label="Answer - Combined Gemini and Mistral")
58
 
59
+ def pdf_qa(file, question):
60
+ if file is None:
61
+ return "Please upload a PDF file first."
62
+ return initialize(file.name, question)
63
 
64
  # Create Gradio Interface
65
  gr.Interface(