Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| from pathlib import Path | |
| from textwrap import dedent | |
| from types import SimpleNamespace | |
| import gradio as gr | |
| from charset_normalizer import detect | |
| from chromadb.config import Settings | |
| from epub2txt import epub2txt | |
| from langchain.chains import RetrievalQA | |
| from langchain.docstore.document import Document | |
| from langchain.document_loaders import ( | |
| CSVLoader, | |
| Docx2txtLoader, | |
| PDFMinerLoader, | |
| TextLoader, | |
| ) | |
| # from constants import CHROMA_SETTINGS, SOURCE_DIRECTORY, PERSIST_DIRECTORY | |
| from langchain.embeddings import HuggingFaceInstructEmbeddings | |
| from langchain.llms import HuggingFacePipeline | |
| from langchain.text_splitter import ( | |
| RecursiveCharacterTextSplitter, | |
| ) | |
| import torch | |
| # FAISS instead of PineCone | |
| from langchain.vectorstores import Chroma | |
| from loguru import logger | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline | |
| import argparse | |
| parser = argparse.ArgumentParser('LocalGPT falcon', add_help=False) | |
| parser.add_argument('--device_type', type=str, default="cpu", choices=["cpu", "mps", "cuda"], help='device type', ) | |
| args = parser.parse_args() | |
| ROOT_DIRECTORY = Path(__file__).parent | |
| PERSIST_DIRECTORY = f"{ROOT_DIRECTORY}/db" | |
| # Define the Chroma settings | |
| CHROMA_SETTINGS = Settings( | |
| chroma_db_impl="duckdb+parquet", | |
| persist_directory=PERSIST_DIRECTORY, | |
| anonymized_telemetry=False, | |
| ) | |
| ns = SimpleNamespace(qa=None) | |
| # INSTRUCTORS_EMBEDDINGS_MODEL = "hkunlp/instructor-xl" | |
| INSTRUCTORS_EMBEDDINGS_MODEL = "hkunlp/instructor-large" | |
| # INSTRUCTORS_EMBEDDINGS_MODEL = "hkunlp/instructor-large" | |
| # INSTRUCTORS_EMBEDDINGS_MODEL = "hkunlp/instructor-base" | |
| def load_single_document(file_path: str or Path) -> Document: | |
| """ingest.py""" | |
| # Loads a single document from a file path | |
| # encoding = detect(open(file_path, "rb").read()).get("encoding", "utf-8") | |
| encoding = detect(Path(file_path).read_bytes()).get("encoding", "utf-8") | |
| if file_path.endswith(".txt"): | |
| if encoding is None: | |
| logger.warning( | |
| f" {file_path}'s encoding is None " | |
| "Something is fishy, return empty str " | |
| ) | |
| return Document(page_content="", metadata={"source": file_path}) | |
| try: | |
| loader = TextLoader(file_path, encoding=encoding) | |
| except Exception as exc: | |
| logger.warning(f" {exc}, return dummy ") | |
| return Document(page_content="", metadata={"source": file_path}) | |
| elif file_path.endswith(".pdf"): | |
| loader = PDFMinerLoader(file_path) | |
| elif file_path.endswith(".csv"): | |
| loader = CSVLoader(file_path) | |
| elif Path(file_path).suffix in [".docx"]: | |
| try: | |
| loader = Docx2txtLoader(file_path) | |
| except Exception as exc: | |
| logger.error(f" {file_path} errors: {exc}") | |
| return Document(page_content="", metadata={"source": file_path}) | |
| elif Path(file_path).suffix in [".epub"]: # for epub? epub2txt unstructured | |
| try: | |
| _ = epub2txt(file_path) | |
| except Exception as exc: | |
| logger.error(f" {file_path} errors: {exc}") | |
| return Document(page_content="", metadata={"source": file_path}) | |
| return Document(page_content=_, metadata={"source": file_path}) | |
| else: | |
| if encoding is None: | |
| logger.warning( | |
| f" {file_path}'s encoding is None " | |
| "Likely binary files, return empty str " | |
| ) | |
| return Document(page_content="", metadata={"source": file_path}) | |
| try: | |
| loader = TextLoader(file_path) | |
| except Exception as exc: | |
| logger.error(f" {exc}, returnning empty string") | |
| return Document(page_content="", metadata={"source": file_path}) | |
| return loader.load()[0] | |
| def greet(name): | |
| """Test.""" | |
| logger.debug(f" name: [{name}] ") | |
| return "Hello " + name + "!!" | |
| def upload_files(files): | |
| """Upload files.""" | |
| try: | |
| file_paths = [file.name for file in files] | |
| except: | |
| file_paths = [files] | |
| logger.info(file_paths) | |
| res = ingest(file_paths) | |
| logger.info("Processed:\n{res}") | |
| del res | |
| ns.qa = load_qa() | |
| return file_paths | |
| def ingest( | |
| file_paths: list | |
| ): | |
| """Gen Chroma db. | |
| torch.cuda.is_available() | |
| file_paths = | |
| [] | |
| """ | |
| logger.info("Doing ingest...") | |
| documents = [] | |
| for file_path in file_paths: | |
| documents.append(load_single_document(f"{file_path}")) | |
| text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) | |
| texts = text_splitter.split_documents(documents) | |
| logger.info(f"Loaded {len(documents)} documents ") | |
| logger.info(f"Split into {len(texts)} chunks of text") | |
| # Create embeddings | |
| logger.info(f"Load InstructEmbeddings model: {INSTRUCTORS_EMBEDDINGS_MODEL}") | |
| embeddings = HuggingFaceInstructEmbeddings( | |
| model_name=INSTRUCTORS_EMBEDDINGS_MODEL, model_kwargs={"device": args.device_type} | |
| ) | |
| db = Chroma.from_documents( | |
| texts, | |
| embeddings, | |
| persist_directory=PERSIST_DIRECTORY, | |
| client_settings=CHROMA_SETTINGS, | |
| ) | |
| db.persist() | |
| db = None | |
| logger.info("Done ingest") | |
| return [ | |
| [Path(doc.metadata.get("source")).name, len(doc.page_content)] | |
| for doc in documents | |
| ] | |
| # https://huggingface.co/tiiuae/falcon-7b-instruct | |
| def gen_local_llm(): | |
| """Gen a local llm. | |
| localgpt run_localgpt | |
| """ | |
| model = "tiiuae/falcon-7b-instruct" | |
| if args.device_type == "cuda": | |
| tokenizer = AutoTokenizer.from_pretrained(model) | |
| else: # cpu | |
| tokenizer=AutoTokenizer.from_pretrained(model) | |
| model=AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True) | |
| pipe = pipeline( | |
| "text-generation", | |
| model=model, | |
| tokenizer=tokenizer, | |
| torch_dtype=torch.float32 if args.device_type =="cpu" else torch.bfloat16, | |
| trust_remote_code=True, | |
| device_map="cpu" if args.device_type =="cpu" else "auto", | |
| max_length=2048, | |
| temperature=0, | |
| top_p=0.95, | |
| top_k=10, | |
| repetition_penalty=1.15, | |
| num_return_sequences=1, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| local_llm = HuggingFacePipeline(pipeline=pipe) | |
| return local_llm | |
| def load_qa(): | |
| """Gen qa.""" | |
| logger.info("Doing qa") | |
| embeddings = HuggingFaceInstructEmbeddings( | |
| model_name=INSTRUCTORS_EMBEDDINGS_MODEL, model_kwargs={"device": args.device_type} | |
| ) | |
| # xl 4.96G, large 3.5G, | |
| db = Chroma( | |
| persist_directory=PERSIST_DIRECTORY, | |
| embedding_function=embeddings, | |
| client_settings=CHROMA_SETTINGS, | |
| ) | |
| retriever = db.as_retriever() | |
| llm = gen_local_llm() # "tiiuae/falcon-7b-instruct" | |
| qa = RetrievalQA.from_chain_type( | |
| llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True | |
| ) | |
| logger.info("Done qa") | |
| return qa | |
| def main1(): | |
| """Lump codes""" | |
| with gr.Blocks() as demo: | |
| iface = gr.Interface(fn=greet, inputs="text", outputs="text") | |
| iface.launch() | |
| demo.launch() | |
| def main(): | |
| """Do blocks.""" | |
| logger.info(f"ROOT_DIRECTORY: {ROOT_DIRECTORY}") | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| with gr.Accordion("Info", open=False): | |
| _ = """ | |
| Talk to your docs (.pdf, .docx, .csv, .txt .md). It | |
| takes quite a while to ingest docs (10-30 min. depending | |
| on net, RAM, CPU etc.). | |
| """ | |
| gr.Markdown(dedent(_)) | |
| title = """ | |
| <div style="text-align: center;"> | |
| <h1>LocalGPT with Falcon</h1> | |
| <p style="text-align: center;">Upload your docs (.pdf, .docx, .csv, .txt .md) by clicking the "Load docs to LangChain" and wait until the upload is complete, <br /> | |
| when everything is ready, you can start asking questions about the docs <br /> | |
| </div> | |
| """ | |
| gr.HTML(title) | |
| with gr.Tab("Upload files"): | |
| # Upload files and generate embeddings database | |
| file_output = gr.File() | |
| upload_button = gr.UploadButton( | |
| "Load docs to LangChain", | |
| file_count="multiple", | |
| ) | |
| upload_button.upload(upload_files, upload_button, file_output) | |
| chatbot = gr.Chatbot() | |
| msg = gr.Textbox(label="Query") | |
| clear = gr.Button("Clear") | |
| def respond(message, chat_history): | |
| if ns.qa is None: # no files processed yet | |
| bot_message = "Provide some file(s) for processsing first." | |
| chat_history.append((message, bot_message)) | |
| return "", chat_history | |
| try: | |
| res = ns.qa(message) | |
| answer, docs = res["result"], res["source_documents"] | |
| bot_message = f"{answer}" | |
| except Exception as exc: | |
| logger.error(exc) | |
| bot_message = f"bummer! {exc}" | |
| chat_history.append((message, bot_message)) | |
| return "", chat_history | |
| msg.submit(respond, [msg, chatbot], [msg, chatbot]) | |
| clear.click(lambda: None, None, chatbot, queue=False) | |
| try: | |
| from google import colab | |
| share = True # start share when in colab | |
| except Exception: | |
| share = False | |
| demo.launch(share=share) | |
| if __name__ == "__main__": | |
| main() | |