arielrocha01/ScanArct

1import google.generativeai as genai2from langchain import PromptTemplate3from langchain.chains.question_answering import load_qa_chain4from langchain.document_loaders import PyPDFDirectoryLoader5from langchain.text_splitter import RecursiveCharacterTextSplitter6from langchain.vectorstores import Chroma7from langchain_google_genai import GoogleGenerativeAIEmbeddings8from langchain_google_genai import ChatGoogleGenerativeAI9import os10import gradio as gr11from PyPDF2 import PdfReader12 13os.environ["GOOGLE_API_KEY"] = "AIzaSyAnEzmacVn7SP2B71ayK3XzFkips7aVNBU"14genai.configure(api_key=os.environ["GOOGLE_API_KEY"])15embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")16model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=1)17 18prompt_template = """19  Com base no artigo disponibilizado para o contexto, responda as perguntas de acordo com esse documento.20\n\n21  Contexto:\n {context}?\n22  Pergunta: \n{question}\n23 24  Resposta:25"""26 27prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])28 29vector_store = None30chain = None31 32def text_to_embedding(data):33    global vector_store34    global chain35 36    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)37    content = "\n\n".join([str(data)])38    texts = text_splitter.split_text(content)39    vector_store = Chroma.from_texts(texts, embeddings).as_retriever()40    chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)41 42def get_text_pdf(file):43    with open(file, 'rb') as pdf_file:44        pdf_reader = PdfReader(pdf_file)45        text = ''46        for page_num in range(len(pdf_reader.pages)):47            page = pdf_reader.pages[page_num]48            text += page.extract_text()49    text_to_embedding(text)50    return text51 52def get_question(question):53    global vector_store54    global chain55 56    docs = vector_store.get_relevant_documents(question)57    58    response = chain(59        {"input_documents": docs, "question": question},60        return_only_outputs=True)61    return response62 63 64with gr.Blocks() as demo:65 66    with gr.Row():67        with gr.Column(scale=1, min_width=600):68            text1 = gr.Interface(fn=get_text_pdf, inputs="file", outputs="text")69            text2 = gr.Interface(fn=get_question, inputs="text", outputs="text")70 71demo.launch()