Roberta2024 commited on
Commit
ae38eb4
1 Parent(s): 99694ff

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +56 -0
app.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.chains.question_answering import load_qa_chain
7
+ import torch
8
+ from transformers import AutoTokenizer, AutoModelForCausalLM
9
+
10
+ # Load Mistral model
11
+ model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
12
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
13
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
14
+ dtype = torch.bfloat16
15
+ model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
16
+
17
+ async def initialize(file_path, question):
18
+ 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 Context: \n {context}?\n Question: \n {question} \n Answer: """
19
+ prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
20
+
21
+ if os.path.exists(file_path):
22
+ pdf_loader = PyPDFLoader(file_path)
23
+ pages = pdf_loader.load_and_split()
24
+ context = "\n".join(str(page.page_content) for page in pages[:30])
25
+
26
+ # Prepare input for Mistral model
27
+ input_text = prompt.format(context=context, question=question)
28
+ inputs = tokenizer.encode(input_text, return_tensors='pt').to(device)
29
+
30
+ # Generate the output
31
+ with torch.no_grad():
32
+ outputs = model.generate(inputs, max_length=500) # Adjust max_length as needed
33
+
34
+ # Decode and return the output
35
+ answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
36
+ return answer
37
+ else:
38
+ return "Error: Unable to process the document. Please ensure the PDF file is valid."
39
+
40
+ # Define Gradio Interface
41
+ input_file = gr.File(label="Upload PDF File")
42
+ input_question = gr.Textbox(label="Ask about the document")
43
+ output_text = gr.Textbox(label="Answer - Mistral Model")
44
+
45
+ async def pdf_qa(file, question):
46
+ answer = await initialize(file.name, question)
47
+ return answer
48
+
49
+ # Create Gradio Interface
50
+ gr.Interface(
51
+ fn=pdf_qa,
52
+ inputs=[input_file, input_question],
53
+ outputs=output_text,
54
+ title="RAG Knowledge Retrieval using Mistral Model",
55
+ description="Upload a PDF file and ask questions about the content."
56
+ ).launch()