teste streamlit

app.py

@@ -1,14 +1,69 @@
-import gradio as gr
-import spaces
 import torch
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 
-@spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+import streamlit as st
+# from streamlit_pdf_viewer import pdf_viewer
+from utils import read_file_pdf, fix_type, extract_content_in_pdf, EXTENSIONS_FILES, EXTENSIONS_IMG_FILES
+from rag_utils import create_split_doc, store_docs, create_rag_chain
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+st.write("## Pergunte qualquer coisa para seu arquivo.")
+st.write(
+    ":dog: Faça o upload do seu arquivo e pergunte qualquer coisa a ele! Este código é open source e disponível [aqui](https://github.com/FelipeErmeson) no GitHub. :grin:"
+)
+st.sidebar.write("## Upload :gear:")
+
+# Increased file size limit
+MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
+
+# UI Layout
+col1, col2 = st.columns(2)
+my_upload = st.sidebar.file_uploader("Upload da imagem", type=["png", "jpg", "jpeg", "pdf"])
+
+# Information about limitations
+with st.sidebar.expander("ℹ️ Diretrizes da Imagem"):
+    st.write("""
+    - Tamanho máximo do arquivo: 10MB
+    - Imagens enormes são automaticamente redimensionadas
+    - Formatos suportados: PNG, JPG, JPEG, PDF
+    - Processamento de tempo depende da GPU alocada
+    """)
+
+# Processa o arquivo
+if my_upload is not None:
+    if my_upload.size > MAX_FILE_SIZE:
+        st.error(f"O arquivo excede o limite. Por favor, realize o upload de um arquivo que contenha no máximo {MAX_FILE_SIZE/1024/1024:.1f}MB.")
+    else:
+        print(my_upload)
+        print(my_upload.type)
+        # binary_data = my_upload.getvalue()
+        # pdf_viewer(input=binary_data, width=700)
+        # read_file_pdf()
+        # fix_image(upload=my_upload)
+
+        file, type_file = fix_type(my_upload)
+        print('type_file', type_file)
+        texto_extraido = None
+        if type_file in EXTENSIONS_FILES:
+            texto_extraido = extract_content_in_pdf(file)
+        elif type_file in EXTENSIONS_IMG_FILES:
+            pass
+        
+        print(texto_extraido)
+
+        if texto_extraido is not None:
+            col1.write("#### Texto extraído:")
+            col1.write(texto_extraido)
+
+            docs_splitted = create_split_doc(texto_extraido)
+            vector_store = store_docs(docs_splitted)
+
+            if question := col2.chat_input("Faça uma pergunta ao seu documento!"):
+                col2.write("📌 " + question)
+
+                rag_chain = create_rag_chain(vector_store)
+                resposta = rag_chain.run(question)
+                col2.write("🎩 " + resposta)
+
+            

rag_utils.py

@@ -0,0 +1,95 @@
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain.llms import HuggingFacePipeline
+
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import RetrievalQA
+
+import torch
+print(torch.cuda.is_available())
+print(torch.cuda.get_device_name(0))
+device = f'cuda:{torch.cuda.current_device()}' if torch.cuda.is_available() else 'cpu'
+
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+cache_dir = "/home/user/.cache/huggingface" #"./model/qwen-awq" #"/home/felipe/.cache/huggingface/transformers" #"/home/user/.cache/huggingface"
+
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+model_name = "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8" #"Qwen/Qwen2.5-7B-Instruct-AWQ" #"Qwen/Qwen2.5-7B-Instruct"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto",
+    trust_remote_code=True,
+    cache_dir=cache_dir
+)
+model.to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, cache_dir=cache_dir)
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=512,
+    temperature=0.1,
+    do_sample=False
+)
+
+# Adapta para LangChain
+llm = HuggingFacePipeline(pipeline=pipe)
+
+def create_split_doc(raw_text):
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    docs = text_splitter.create_documents([raw_text])
+
+    return docs
+
+def store_docs(docs):
+    vectorstore = FAISS.from_documents(docs, embedding_model)
+    return vectorstore
+
+def create_template():
+    prompt_template = PromptTemplate(
+        input_variables=["context", "question"],
+        template="""
+Você é um especialista em extrair informações em documentos.
+Com base nas informações a seguir, forneça a melhor resposta.
+Caso não tenha certeza da resposta, prefira falar que não sabe responder tal pergunta.
+Responda de maneira amigável e clara.
+
+Contexto:
+{context}
+
+Pergunta:
+{question}
+"""
+)
+    return prompt_template
+
+def create_rag_chain(vectorstore):
+    rag_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    retriever=vectorstore.as_retriever(),
+    chain_type="stuff",
+    chain_type_kwargs={"prompt": create_template()}
+)
+    return rag_chain
+
+
+if __name__ == '__main__':
+    pass
+
+# resposta = rag_chain.run(pergunta)
+
+
+
+# pergunta = "Qual o número da nfse?"
+# resposta = rag_chain.run(pergunta)
+
+# print("📌 Pergunta:", pergunta)
+# print("🎩 Resposta do Analista Fiscal:\\n", resposta)

utils.py

@@ -0,0 +1,69 @@
+from PyPDF2 import PdfReader
+from PIL import Image
+from io import BytesIO
+
+EXTENSIONS_IMG_FILES = ['jpeg', 'jpg', 'png']
+EXTENSIONS_FILES = ['pdf']
+EXTENSIONS_ALLOWED = EXTENSIONS_IMG_FILES + EXTENSIONS_FILES
+
+# Max dimensions for processing
+MAX_IMAGE_SIZE = 2000  # pixels
+
+def fix_type(file_upload):
+    if isinstance(file_upload, str):
+        print('teste: str')
+    else:
+        type_file = file_upload.type.split('/')[-1]
+        if type_file in EXTENSIONS_IMG_FILES:
+            return read_file_img(file_upload), type_file
+        elif type_file in EXTENSIONS_FILES:
+            return read_file_pdf(file_upload), type_file
+
+# Resize image while maintaining aspect ratio
+def resize_image(image, max_size):
+    width, height = image.size
+    if width <= max_size and height <= max_size:
+        return image
+    
+    if width > height:
+        new_width = max_size
+        new_height = int(height * (max_size / width))
+    else:
+        new_height = max_size
+        new_width = int(width * (max_size / height))
+    
+    return image.resize((new_width, new_height), Image.LANCZOS)
+
+def process_image(image_bytes):
+    try:
+        image = Image.open(BytesIO(image_bytes))
+        # Resize large images to prevent memory issues
+        # resized = resize_image(image, MAX_IMAGE_SIZE)
+        return image
+    except Exception as e:
+        # st.error(f"Error processing image: {str(e)}")
+        return None
+
+def read_file_img(file_img):
+    image_bytes = file_img.getvalue()
+    img_pil = process_image(image_bytes)
+    return img_pil
+
+def read_file_pdf(file_pdf):
+    # image_bytes = file_pdf.getvalue()
+    reader = PdfReader(file_pdf)
+    return reader
+
+def extract_content_in_pdf(reader):
+    raw_text = ""
+    for page in reader.pages:
+        text = page.extract_text()
+        if text:
+            raw_text += text + "\\n"
+    
+    return raw_text
+
+# st.write(f"O PDF tem {num_pages} páginas.")
+
+if __name__ == '__main__':
+    pass