Initial commit

Dockerfile

@@ -0,0 +1,25 @@
+FROM debian
+# FROM python:3.10.12
+
+RUN apt-get update && apt-get upgrade
+
+RUN apt-get install python3.9
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Assegura que huggingface-cli está instalado
+RUN pip install huggingface-hub
+
+COPY ./app.py /code/app.py
+
+COPY ./check_token.py /code/check_token.py
+
+COPY ./inference.py /code/inference.py
+
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+ENTRYPOINT ["/entrypoint.sh"] && /bin/bash

app.py

@@ -0,0 +1,146 @@
+import os
+import shutil
+import shutil
+import numpy as np
+from uuid import uuid4
+from io import BytesIO
+from pydantic import BaseModel
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, UploadFile, File, Response, status
+from llama_index.readers import StringIterableReader, PDFReader, SimpleDirectoryReader
+from llama_index import (
+    VectorStoreIndex,
+    ServiceContext,
+    set_global_service_context,
+)
+from pyngrok import ngrok
+import inference
+
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class Message(BaseModel):
+  content: str
+
+if not os.path.exists("tmp"):
+  os.mkdir("tmp")
+
+vector_stores = {}
+
+@app.post("/retriveal/ingest")
+async def store_file(
+    file: UploadFile = File(...)
+):
+
+  try:
+
+    print(file.filename)
+    id = str(uuid4())
+    file_location = f"tmp/{id}"
+
+    if not os.path.exists(file_location):
+      os.mkdir(file_location)
+
+    with open(f"{file_location}/{file.filename}",  "wb+") as f:
+      shutil.copyfileobj(file.file, f)
+
+    pdf = SimpleDirectoryReader(f"tmp/{id}").load_data()
+
+    vector_stores[id] = VectorStoreIndex.from_documents(pdf)
+
+    return jsonable_encoder({"uuid": id})
+
+  except Exception as e:
+
+    # response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
+    return jsonable_encoder({"error": str(e)})
+
+@app.post("/retriveal/ingest/{id}")
+async def store_file_with_id(
+    id,
+    file: UploadFile = File(...)
+):
+
+  try:
+
+    print(file.filename)
+
+    if(id == None or id == ""):
+      raise Exception("Id is required")
+
+    file_location = f"tmp/{id}"
+
+    if not os.path.exists(file_location):
+      os.mkdir(file_location)
+
+    with open(f"{file_location}/{file.filename}",  "wb+") as f:
+      shutil.copyfileobj(file.file, f)
+
+    pdf = SimpleDirectoryReader(f"tmp/{id}").load_data()
+
+    vector_stores[id] = VectorStoreIndex.from_documents(pdf)
+
+    return jsonable_encoder({"uuid": id})
+
+  except Exception as e:
+
+    # response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
+    return jsonable_encoder({"error": str(e)})
+
+@app.delete("/session/{id}")
+async def delete_session(id):
+  try:
+    shutil.rmtree(f"tmp/{id}")
+    return jsonable_encoder({"message": "ok"})
+  except Exception as e:
+    return jsonable_encoder({"error": str(e)})
+
+@app.post("/retriveal/{id}")
+async def inference(
+    id,
+    message: Message
+):
+
+    if(id == None or id == ""):
+      raise Exception("Id is required")
+
+    query = message.content
+
+    query_engine = vector_stores[id].as_query_engine()
+
+    inference = query_engine.query(query)
+
+    return inference
+
+
+def stream_inference(gen):
+  for token in gen:
+    yield token
+
+
+@app.post("/retriveal/stream/{id}")
+async def inference(
+    id,
+    message: Message
+):
+
+    if(id == None or id == ""):
+      raise Exception("Id is required")
+
+    query = message.content
+
+    query_engine = vector_stores[id].as_query_engine(streaming=True)
+
+    gen = query_engine.query(query).response_gen
+
+    return StreamingResponse(stream_inference(gen))

check_token.py

@@ -0,0 +1,6 @@
+import os
+from huggingface_hub._login import _login
+
+HF_KEY = os.environ["HF_KEY"]
+
+_login(token=HF_KEY, add_to_git_credential=False)

docker

@@ -0,0 +1 @@
+failed to get console mode for stdout: Identificador inválido.

entrypoint.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+# Verifique o token
+python /code/check_token.py $HF_TOKEN
+
+# Verifique o status de saída
+if [ $? -eq 0 ]; then
+    echo "Login bem-sucedido!"
+else
+    echo "Falha no login."
+    exit 1
+fi
+
+# Inicie o seu aplicativo
+exec "$@"

inference.py

@@ -0,0 +1,59 @@
+import logging
+import sys
+from IPython.display import Markdown, display
+import torch
+from llama_index.llms.huggingface import HuggingFaceLLM
+from llama_index.prompts import PromptTemplate
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+from llama_index import (
+    ServiceContext,
+    set_global_service_context,
+)
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
+
+
+# Model names (make sure you have access on HF)
+LLAMA2_7B = "meta-llama/Llama-2-7b-hf"
+LLAMA2_7B_CHAT = "meta-llama/Llama-2-7b-chat-hf"
+LLAMA2_13B = "meta-llama/Llama-2-13b-hf"
+LLAMA2_13B_CHAT = "meta-llama/Llama-2-13b-chat-hf"
+LLAMA2_70B = "meta-llama/Llama-2-70b-hf"
+LLAMA2_70B_CHAT = "meta-llama/Llama-2-70b-chat-hf"
+
+selected_model = LLAMA2_7B_CHAT
+
+SYSTEM_PROMPT = """Você é um assistente de IA que responde a perguntas de maneira amigável, com base nos documentos fornecidos. Aqui estão algumas regras que você sempre segue:
+- Gerar saídas legíveis para humanos, evitando criar texto sem sentido.
+- Gerar apenas a saída solicitada, sem incluir qualquer outro idioma antes ou depois da saída solicitada.
+- Nunca agradecer, expressar felicidade em ajudar, mencionar que é um agente de IA, etc. Apenas responda diretamente.
+- Gerar linguagem profissional geralmente usada em documentos comerciais na América do Norte.
+- Nunca gerar linguagem ofensiva ou obscena.
+- Traduza as suas respostas sempre para Português Brasileiro. Nunca responsa nada em inglês.
+"""
+
+query_wrapper_prompt = PromptTemplate(
+    "[INST]<<SYS>>\n" + SYSTEM_PROMPT + "<</SYS>>\n\n{query_str}[/INST] "
+)
+
+llm = HuggingFaceLLM(
+    context_window=4096,
+    max_new_tokens=2048,
+    generate_kwargs={"temperature": 0.0, "do_sample": False},
+    query_wrapper_prompt=query_wrapper_prompt,
+    tokenizer_name=selected_model,
+    model_name=selected_model,
+    device_map="auto",
+    # change these settings below depending on your GPU
+    # model_kwargs={"torch_dtype": torch.float16, "load_in_8bit": True},
+)
+
+
+embed_model = HuggingFaceEmbedding(model_name="neuralmind/bert-base-portuguese-cased")
+# embed_model = FlagModel("BAAI/bge-m3")
+
+service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
+
+set_global_service_context(service_context)

requirements.txt

@@ -0,0 +1,13 @@
+langchain==0.1.7 
+fastapi==0.109.2 
+python-multipart==0.0.9 
+uvicorn==0.27.1 
+pypdf==4.0.1
+pyngrok==7.1.2
+ipywidgets==7.7.1 
+bitsandbytes==0.42.0
+accelerate==0.27.2
+llama-index==0.9.47
+torch==2.0.1
+transformers==4.37.2
+pyngrok