sidcww commited on
Commit
60a09fb
·
verified ·
1 Parent(s): 8e4c5ee

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +78 -0
app.py ADDED
@@ -0,0 +1,78 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from langsmith import LangSmithClient # Hypothetical import for LangSmith
11
+
12
+ # Configure Gemini API
13
+ genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
14
+
15
+ # Load Mistral model
16
+ model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
17
+ mistral_tokenizer = AutoTokenizer.from_pretrained(model_path)
18
+ 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
+ # Initialize LangSmith client
23
+ langsmith_client = LangSmithClient(api_key=os.getenv("LANGSMITH_API_KEY"))
24
+
25
+ def initialize(file_path, question):
26
+ try:
27
+ model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
28
+ prompt_template = """Answer the question as precise as possible using the provided context. If the answer is
29
+ not contained in the context, say "answer not available in context" \n\n
30
+ Context: \n {context}?\n
31
+ Question: \n {question} \n
32
+ Answer:
33
+ """
34
+ prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
35
+
36
+ if os.path.exists(file_path):
37
+ pdf_loader = PyPDFLoader(file_path)
38
+ pages = pdf_loader.load_and_split()
39
+ context = "\n".join(str(page.page_content) for page in pages[:10])
40
+ stuff_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
41
+ stuff_answer = stuff_chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
42
+ gemini_answer = stuff_answer['output_text']
43
+
44
+ # Use Mistral model for additional text generation
45
+ mistral_prompt = f"Based on this answer: {gemini_answer}\nGenerate a follow-up question:"
46
+ mistral_inputs = mistral_tokenizer.encode(mistral_prompt, return_tensors='pt').to(device)
47
+ with torch.no_grad():
48
+ mistral_outputs = mistral_model.generate(mistral_inputs, max_length=100)
49
+ mistral_output = mistral_tokenizer.decode(mistral_outputs[0], skip_special_tokens=True)
50
+
51
+ combined_output = f"Gemini Answer: {gemini_answer}\n\nMistral Follow-up: {mistral_output}"
52
+
53
+ # Log and evaluate the results using LangSmith
54
+ evaluation = langsmith_client.evaluate(gemini_answer, mistral_output, question)
55
+ return combined_output + f"\n\nEvaluation: {evaluation}"
56
+ else:
57
+ return "Error: Unable to process the document. Please ensure the PDF file is valid."
58
+ except Exception as e:
59
+ return f"An error occurred: {str(e)}"
60
+
61
+ # Define Gradio Interface
62
+ input_file = gr.File(label="Upload PDF File")
63
+ input_question = gr.Textbox(label="Ask about the document")
64
+ output_text = gr.Textbox(label="Answer - Combined Gemini and Mistral")
65
+
66
+ def pdf_qa(file, question):
67
+ if file is None:
68
+ return "Please upload a PDF file first."
69
+ return initialize(file.name, question)
70
+
71
+ # Create Gradio Interface
72
+ gr.Interface(
73
+ fn=pdf_qa,
74
+ inputs=[input_file, input_question],
75
+ outputs=output_text,
76
+ title="RAG Knowledge Retrieval using Gemini API and Mistral Model",
77
+ description="Upload a PDF file and ask questions about the content."
78
+ ).launch()