Upload 7 files

app.py

@@ -0,0 +1,64 @@
+import fn
+import gradio as gr
+from refresh import create_refresh_button
+
+def fn_search(dir, query):
+    args = {
+        'dir': dir,
+        'query': query,
+        'k': 1,
+    }
+    vector_store = fn.vector_load(args)
+    result, detail = fn.search(vector_store, args)
+
+    return result
+
+with gr.Blocks() as demo:
+    title = gr.Markdown('# FAISS(Naive RAG)')
+    info = gr.Markdown()
+    upload_dir = gr.Dropdown(
+        choices=fn.load_dirs(),
+        label='name',
+        show_label=True,
+        interactive=True,
+        allow_custom_value=True,
+    )
+
+    with gr.Tab('Upload'):
+        chunk_size = gr.Textbox(
+            value=0,
+            label='(optional) chunk_size if split',
+            show_label=True,
+            interactive=True,
+        )
+        rag_zip = gr.UploadButton(
+            label='Zip Upload and save rag_dir',
+            interactive=True,
+        )
+        create_refresh_button(gr, upload_dir, lambda: None, lambda: {'choices': fn.load_dirs()}, 'refresh-button', interactive=True)
+
+    with gr.Tab('Search'):
+        query = gr.Textbox(
+            lines=1,
+            label='query',
+            show_label=True,
+            interactive=True,
+            show_copy_button=True,
+            )
+        search_button = gr.Button(value='search')
+        result = gr.Textbox(label='result', show_label=True, show_copy_button=True)
+
+    search_button.click(
+        fn=fn_search,
+        inputs=[upload_dir, query],
+        outputs=[result],
+        )
+        
+    rag_zip.upload(
+        fn=fn.upload,
+        inputs=[upload_dir, chunk_size, rag_zip],
+        outputs=[info],
+        )
+
+if __name__ == '__main__':
+    demo.launch()

fn.py

@@ -0,0 +1,120 @@
+import os
+import shutil
+import tempfile
+import requests
+
+from langchain.document_loaders.generic import GenericLoader
+from langchain.document_loaders.blob_loaders import FileSystemBlobLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.vectorstores.utils import DistanceStrategy
+from langchain.text_splitter import CharacterTextSplitter
+
+os.makedirs('store', exist_ok = True)
+
+def download(args: dict):
+    if not 'dir' in args:
+        raise ValueError('require dir')
+
+    if 'zip_url' in args:
+        res = requests.get(args['zip_url'])
+
+        with tempfile.NamedTemporaryFile(suffix=".zip") as t:
+            with open(t.name, 'wb') as f:
+                f.write(res.content)
+
+            if os.path.exists(f"store/{args['dir']}"):
+                shutil.rmtree(f"store/{args['dir']}")
+            shutil.unpack_archive(t.name, f"store/{args['dir']}")
+    elif 'url' in args:
+        os.makedirs(f"store/{args['dir']}", exist_ok=True)
+        res = requests.get(args['url'])
+
+        filepath = f"store/{args['dir']}/{os.path.basename(args['url'])}"
+        with open(filepath, 'wb') as f:
+            f.write(res.content)
+    elif 'text' in args:
+        os.makedirs(f"store/{args['dir']}", exist_ok=True)
+
+        filepath = f"store/{args['dir']}/text.txt"
+        with open(filepath, 'w', encoding='utf-8') as f:
+            f.write(args['text'])
+
+def docs_load(args: dict):
+    loader = GenericLoader.from_filesystem(
+        path=f"store/{args['dir']}",
+        glob="**/[!.]*",
+        show_progress=True,
+    )
+
+    docs = loader.load()
+    return docs
+
+def chunk_split(docs, chunk_size):
+    text_splitter = CharacterTextSplitter(
+        separator='\n\n',
+        chunk_size=chunk_size,
+        chunk_overlap=0,
+        length_function=len
+    )
+    chunk_docs = text_splitter.create_documents([doc.page_content for doc in docs])
+    return chunk_docs
+
+def vector(docs, args: dict):
+    embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
+    vector_store = FAISS.from_documents(documents=docs,
+                                        embedding=embeddings, 
+                                        distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
+                                        normalize_L2=True)
+    return vector_store
+
+def vector_save(docs, args: dict):
+    vector_store = vector(docs, args)
+    folder_path = f"store/{args['dir']}/vector"
+    vector_store.save_local(folder_path=folder_path)
+
+    return vector_store
+
+def vector_load(args: dict):
+    folder_path = f"store/{args['dir']}/vector"
+    if not os.path.exists(folder_path):
+        raise ValueError(f"missing store/{args['dir']}/vector")
+
+    embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
+    vector_store = FAISS.load_local(folder_path=folder_path,
+                                    embeddings=embeddings, 
+                                    distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
+                                    normalize_L2=True)
+    return vector_store
+
+def search(vector_store, args: dict):
+    results = vector_store.similarity_search_with_score(query=args['query'], k=args['k'])
+    detail = []
+    for r in results:
+        detail.append([r[0].page_content, float(r[1])])
+    return results[0][0].page_content, detail
+
+def load_dirs():
+    dirs = []
+    for name in os.listdir('store'):
+        dirs.append(name)
+
+    return dirs
+
+def upload(dir, chunk_size, file):
+    if not dir:
+        raise ValueError('require dir')
+
+    args = {
+        'dir': dir,
+        'chunk_size': int(chunk_size),
+        }
+    if os.path.exists(f"store/{args['dir']}"):
+        shutil.rmtree(f"store/{args['dir']}")
+    shutil.unpack_archive(file.name, f"store/{args['dir']}")
+    docs = docs_load(args)
+    if args['chunk_size'] > 0:
+        docs = chunk_split(docs, int(chunk_size))
+    vector_save(docs, args)
+
+    return f"saved store/{args['dir']}"

install.bat

@@ -0,0 +1,56 @@
+@echo off
+
+rem -------------------------------------------
+rem NOT guaranteed to work on Windows
+
+set REPOS=https://huggingface.co/spaces/aka7774/faiss
+set APPDIR=faiss
+set VENV=venv
+
+rem -------------------------------------------
+
+set INSTALL_DIR=%~dp0
+cd /d %INSTALL_DIR%
+
+:git_clone
+set DL_URL=%REPOS%
+set DL_DST=%APPDIR%
+git clone %DL_URL% %APPDIR%
+if exist %DL_DST% goto install_python
+
+set DL_URL=https://github.com/git-for-windows/git/releases/download/v2.41.0.windows.3/PortableGit-2.41.0.3-64-bit.7z.exe
+set DL_DST=PortableGit-2.41.0.3-64-bit.7z.exe
+curl -L -o %DL_DST% %DL_URL%
+if not exist %DL_DST% bitsadmin /transfer dl %DL_URL% %DL_DST%
+%DL_DST% -y
+del %DL_DST%
+
+set GIT=%INSTALL_DIR%PortableGit\bin\git
+%GIT% clone %REPOS%
+
+:install_python
+set DL_URL=https://github.com/indygreg/python-build-standalone/releases/download/20240107/cpython-3.10.13+20240107-i686-pc-windows-msvc-shared-install_only.tar.gz
+set DL_DST="%INSTALL_DIR%python.tar.gz"
+curl -L -o %DL_DST% %DL_URL%
+if not exist %DL_DST% bitsadmin /transfer dl %DL_URL% %DL_DST%
+tar -xzf %DL_DST%
+
+set PYTHON=%INSTALL_DIR%python\python.exe
+set PATH=%PATH%;%INSTALL_DIR%python310\Scripts
+
+:install_venv
+cd %APPDIR%
+%PYTHON% -m venv %VENV%
+set PYTHON=%VENV%\Scripts\python.exe
+
+:install_pip
+set DL_URL=https://bootstrap.pypa.io/get-pip.py
+set DL_DST=%INSTALL_DIR%get-pip.py
+curl -o %DL_DST% %DL_URL%
+if not exist %DL_DST% bitsadmin /transfer dl %DL_URL% %DL_DST%
+%PYTHON% %DL_DST%
+
+%PYTHON% -m pip install gradio
+%PYTHON% -m pip install -r requirements.txt
+
+pause

main.py

@@ -0,0 +1,43 @@
+import os
+import sys
+import time
+import signal
+import io
+
+from fastapi import FastAPI, Request, status, Form, UploadFile
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+
+import fn
+import gradio as gr
+from app import demo
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=['*'],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+gr.mount_gradio_app(app, demo, path="/gradio")
+
+@app.post("/upload")
+async def api_save(args: dict):
+    fn.download(args)
+    docs = fn.docs_load(args)
+    if 'chunk_size' in args and args['chunk_size']:
+        docs = fn.chunk_split(docs, int(args['chunk_size']))
+    fn.vector_save(docs, args)
+    return {"status": 0, "result": "saved."}
+
+@app.post("/search")
+async def api_search(args: dict):
+    vector_store = fn.vector_load(args)
+    result, detail = fn.search(vector_store, args)
+    return {"status": 0, "result": result, "detail": detail}

refresh.py

@@ -0,0 +1,23 @@
+
+def create_refresh_button(gr, refresh_component, refresh_method, refreshed_args, elem_class, interactive=True):
+    """
+    Copied from https://github.com/AUTOMATIC1111/stable-diffusion-webui
+    """
+    refresh_symbol = '🔄'
+    def refresh():
+        refresh_method()
+        args = refreshed_args() if callable(refreshed_args) else refreshed_args
+
+        for k, v in args.items():
+            setattr(refresh_component, k, v)
+
+        return gr.update(**(args or {}))
+
+    refresh_button = gr.Button(refresh_symbol, elem_classes=elem_class, interactive=interactive)
+    refresh_button.click(
+        fn=refresh,
+        inputs=[],
+        outputs=[refresh_component]
+    )
+
+    return refresh_button

requirements.txt

@@ -0,0 +1,6 @@
+fastapi
+uvicorn
+faiss-gpu
+langchain
+sentence-transformers
+python-multipart

venv.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/bash
+
+python3 -m venv venv
+curl -kL https://bootstrap.pypa.io/get-pip.py | venv/bin/python
+
+venv/bin/python -m pip install gradio
+venv/bin/python -m pip install -r requirements.txt