Spencer525 commited on
Commit
eefc898
1 Parent(s): 16514f2

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
+ from langchain_core.prompts import PromptTemplate
4
+ from langchain_community.document_loaders import PyPDFLoader
5
+ from langchain_google_genai import ChatGoogleGenerativeAI
6
+ import google.generativeai as genai
7
+ from langchain.chains.question_answering import load_qa_chain
8
+ import torch
9
+ from transformers import AutoTokenizer, AutoModelForCausalLM
10
+
11
+ # Configure Gemini API
12
+ genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
13
+
14
+ # Load Mistral model
15
+ model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
16
+ mistral_tokenizer = AutoTokenizer.from_pretrained(model_path)
17
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
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(
66
+ fn=pdf_qa,
67
+ inputs=[input_file, input_question],
68
+ outputs=output_text,
69
+ title="RAG Knowledge Retrieval using Gemini API and Mistral Model",
70
+ description="Upload a PDF file and ask questions about the content."
71
+ ).launch()