File size: 4,897 Bytes
edc070f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
from fastapi import FastAPI

# from transformers import pipeline
from txtai.embeddings import Embeddings
from txtai.pipeline import Extractor
from langchain.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

import pandas as pd
import sqlite3
import os

# NOTE - we configure docs_url to serve the interactive Docs at the root path
# of the app. This way, we can use the docs as a landing page for the app on Spaces.
app = FastAPI(docs_url="/")
# app = FastAPI()

# pipe = pipeline("text2text-generation", model="google/flan-t5-small")


# @app.get("/generate")
# def generate(text: str):
#     """
#     Using the text2text-generation pipeline from `transformers`, generate text
#     from the given input text. The model used is `google/flan-t5-small`, which
#     can be found [here](https://huggingface.co/google/flan-t5-small).
#     """
#     output = pipe(text)
#     return {"output": output[0]["generated_text"]}


def load_embeddings(
    domain: str = "",
    db_present: bool = True,
    path: str = "sentence-transformers/all-MiniLM-L6-v2",
    index_name: str = "index",
):
    # Create embeddings model with content support
    embeddings = Embeddings({"path": path, "content": True})

    # if Vector DB is not present
    if not db_present:
        return embeddings
    else:
        if domain == "":
            embeddings.load(index_name)  # change this later
        else:
            print(3)
            embeddings.load(f"{index_name}/{domain}")
        return embeddings


def _check_if_db_exists(db_path: str) -> bool:
    return os.path.exists(db_path)


def _text_splitter(doc):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        length_function=len,
    )
    return text_splitter.transform_documents(doc)


def _load_docs(path: str):
    load_doc = WebBaseLoader(path).load()
    doc = _text_splitter(load_doc)
    return doc


def _stream(dataset, limit, index: int = 0):
    for row in dataset:
        yield (index, row.page_content, None)
        index += 1

        if index >= limit:
            break


def _max_index_id(path):
    db = sqlite3.connect(path)

    table = "sections"
    df = pd.read_sql_query(f"select * from {table}", db)
    return {"max_index": df["indexid"].max()}


def _upsert_docs(doc, embeddings, vector_doc_path: str, db_present: bool):
    print(vector_doc_path)
    if db_present:
        print(1)
        max_index = _max_index_id(f"{vector_doc_path}/documents")
        print(max_index)
        embeddings.upsert(_stream(doc, 500, max_index["max_index"]))
        print("Embeddings done!!")
        embeddings.save(vector_doc_path)
        print("Embeddings done - 1!!")
    else:
        print(2)
        embeddings.index(_stream(doc, 500, 0))
        embeddings.save(vector_doc_path)
        max_index = _max_index_id(f"{vector_doc_path}/documents")
        print(max_index)
    # check
    # max_index = _max_index_id(f"{vector_doc_path}/documents")
    # print(max_index)
    return max_index


# def prompt(question):
#     return f"""Answer the following question using only the context below. Say 'no answer' when the question can't be answered.
#             Question: {question}
#             Context: """


# def search(query, question=None):
#     # Default question to query if empty
#     if not question:
#         question = query

#     return extractor([("answer", query, prompt(question), False)])[0][1]


# @app.get("/rag")
# def rag(question: str):
#     # question = "what is the document about?"
#     answer = search(question)
#     # print(question, answer)
#     return {answer}


# @app.get("/index")
# def get_url_file_path(url_path: str):
#     embeddings = load_embeddings()
#     doc = _load_docs(url_path)
#     embeddings, max_index = _upsert_docs(doc, embeddings)
#     return max_index


@app.get("/index/{domain}/")
def get_domain_file_path(domain: str, file_path: str):
    print(domain, file_path)
    print(os.getcwd())
    bool_value = _check_if_db_exists(db_path=f"{os.getcwd()}\index\{domain}\documents")
    print(bool_value)
    if bool_value:
        embeddings = load_embeddings(domain=domain, db_present=bool_value)
        print(embeddings)
        doc = _load_docs(file_path)
        max_index = _upsert_docs(
            doc=doc,
            embeddings=embeddings,
            vector_doc_path=f"index/{domain}",
            db_present=bool_value,
        )
        # print("-------")
    else:
        embeddings = load_embeddings(domain=domain, db_present=bool_value)
        doc = _load_docs(file_path)
        max_index = _upsert_docs(
            doc=doc,
            embeddings=embeddings,
            vector_doc_path=f"index/{domain}",
            db_present=bool_value,
        )
    # print("Final - output : ", max_index)
    return "Executed Successfully!!"