mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-02 17:46:42 +00:00
[Feature] Add document retrieval QA (#5020)
* add langchain * add langchain * Add files via upload * add langchain * fix style * fix style: remove extra space * add pytest; modified retriever * add pytest; modified retriever * add tests to build_on_pr.yml * fix build_on_pr.yml * fix build on pr; fix environ vars * seperate unit tests for colossalqa from build from pr * fix container setting; fix environ vars * commented dev code * add incremental update * remove stale code * fix style * change to sha3 224 * fix retriever; fix style; add unit test for document loader * fix ci workflow config * fix ci workflow config * add set cuda visible device script in ci * fix doc string * fix style; update readme; refactored * add force log info * change build on pr, ignore colossalqa * fix docstring, captitalize all initial letters * fix indexing; fix text-splitter * remove debug code, update reference * reset previous commit * update LICENSE update README add key-value mode, fix bugs * add files back * revert force push * remove junk file * add test files * fix retriever bug, add intent classification * change conversation chain design * rewrite prompt and conversation chain * add ui v1 * ui v1 * fix atavar * add header * Refactor the RAG Code and support Pangu * Refactor the ColossalQA chain to Object-Oriented Programming and the UI demo. * resolved conversation. tested scripts under examples. web demo still buggy * fix ci tests * Some modifications to add ChatGPT api * modify llm.py and remove unnecessary files * Delete applications/ColossalQA/examples/ui/test_frontend_input.json * Remove OpenAI api key * add colossalqa * move files * move files * move files * move files * fix style * Add Readme and fix some bugs. * Add something to readme and modify some code * modify a directory name for clarity * remove redundant directory * Correct a type in llm.py * fix AI prefix * fix test_memory.py * fix conversation * fix some erros and typos * Fix a missing import in RAG_ChatBot.py * add colossalcloud LLM wrapper, correct issues in code review --------- Co-authored-by: YeAnbang <anbangy2@outlook.com> Co-authored-by: Orion-Zheng <zheng_zian@u.nus.edu> Co-authored-by: Zian(Andy) Zheng <62330719+Orion-Zheng@users.noreply.github.com> Co-authored-by: Orion-Zheng <zhengzian@u.nus.edu>
This commit is contained in:
YeAnbang
GitHub
@@ -0,0 +1,119 @@
|
||||
'''
|
||||
Class for loading table type data. please refer to Pandas-Input/Output for file format details.
|
||||
'''
|
||||
|
||||
|
||||
import os
|
||||
import glob
|
||||
import pandas as pd
|
||||
from sqlalchemy import create_engine
|
||||
from colossalqa.utils import drop_table
|
||||
from colossalqa.mylogging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
SUPPORTED_DATA_FORMAT = ['.csv','.xlsx', '.xls','.json','.html','.h5', '.hdf5','.parquet','.feather','.dta']
|
||||
|
||||
class TableLoader:
|
||||
'''
|
||||
Load tables from different files and serve a sql database for database operations
|
||||
'''
|
||||
def __init__(self, files: str,
|
||||
sql_path:str='sqlite:///mydatabase.db',
|
||||
verbose=False, **kwargs) -> None:
|
||||
'''
|
||||
Args:
|
||||
files: list of files (list[file path, name])
|
||||
sql_path: how to serve the sql database
|
||||
**kwargs: keyword type arguments, useful for certain document types
|
||||
'''
|
||||
self.data = {}
|
||||
self.verbose = verbose
|
||||
self.sql_path = sql_path
|
||||
self.kwargs = kwargs
|
||||
self.sql_engine = create_engine(self.sql_path)
|
||||
drop_table(self.sql_engine)
|
||||
|
||||
self.sql_engine = create_engine(self.sql_path)
|
||||
for item in files:
|
||||
path = item[0]
|
||||
dataset_name = item[1]
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"{path} doesn't exists")
|
||||
if not any([path.endswith(i) for i in SUPPORTED_DATA_FORMAT]):
|
||||
raise TypeError(f"{path} not supported. Supported type {SUPPORTED_DATA_FORMAT}")
|
||||
|
||||
logger.info("loading data", verbose=self.verbose)
|
||||
self.load_data(path)
|
||||
logger.info("data loaded", verbose=self.verbose)
|
||||
self.to_sql(path, dataset_name)
|
||||
|
||||
def load_data(self, path):
|
||||
'''
|
||||
Load data and serve the data as sql database.
|
||||
Data must be in pandas format
|
||||
'''
|
||||
files = []
|
||||
# Handle glob expression
|
||||
try:
|
||||
files = glob.glob(path)
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
if len(files)==0:
|
||||
raise ValueError("Unsupported file/directory format. For directories, please use glob expression")
|
||||
elif len(files)==1:
|
||||
path = files[0]
|
||||
else:
|
||||
for file in files:
|
||||
self.load_data(file)
|
||||
|
||||
if path.endswith('.csv'):
|
||||
# Load csv
|
||||
self.data[path] = pd.read_csv(path)
|
||||
elif path.endswith('.xlsx') or path.endswith('.xls'):
|
||||
# Load excel
|
||||
self.data[path] = pd.read_excel(path) # You can adjust the sheet_name as needed
|
||||
elif path.endswith('.json'):
|
||||
# Load json
|
||||
self.data[path] = pd.read_json(path)
|
||||
elif path.endswith('.html'):
|
||||
# Load html
|
||||
html_tables = pd.read_html(path)
|
||||
# Choose the desired table from the list of DataFrame objects
|
||||
self.data[path] = html_tables[0] # You may need to adjust this index
|
||||
elif path.endswith('.h5') or path.endswith('.hdf5'):
|
||||
# Load h5
|
||||
self.data[path] = pd.read_hdf(path, key=self.kwargs.get('key', 'data')) # You can adjust the key as needed
|
||||
elif path.endswith('.parquet'):
|
||||
# Load parquet
|
||||
self.data[path] = pd.read_parquet(path, engine='fastparquet')
|
||||
elif path.endswith('.feather'):
|
||||
# Load feather
|
||||
self.data[path] = pd.read_feather(path)
|
||||
elif path.endswith('.dta'):
|
||||
# Load dta
|
||||
self.data[path] = pd.read_stata(path)
|
||||
else:
|
||||
raise ValueError("Unsupported file format")
|
||||
|
||||
def to_sql(self, path, table_name):
|
||||
'''
|
||||
Serve the data as sql database.
|
||||
'''
|
||||
self.data[path].to_sql(table_name, con=self.sql_engine, if_exists='replace', index=False)
|
||||
logger.info(f"Loaded to Sqlite3\nPath: {path}", verbose=self.verbose)
|
||||
return self.sql_path
|
||||
|
||||
def get_sql_path(self):
|
||||
return self.sql_path
|
||||
|
||||
def __del__(self):
|
||||
if self.sql_engine:
|
||||
drop_table(self.sql_engine)
|
||||
self.sql_engine.dispose()
|
||||
del self.data
|
||||
del self.sql_engine
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user