Spaces:
Build error
Build error
add intention mapper to embedding flow , refactor into classes
Browse files- app.py +102 -12
- requirements.txt +9 -1
- utils.py +17 -15
app.py
CHANGED
@@ -18,7 +18,12 @@ from globalvars import API_BASE, intention_prompt, tasks
|
|
18 |
from dotenv import load_dotenv
|
19 |
import re
|
20 |
from utils import load_env_variables
|
21 |
-
|
|
|
|
|
|
|
|
|
|
|
22 |
|
23 |
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'
|
24 |
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
|
@@ -29,23 +34,108 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
29 |
|
30 |
hf_token, yi_token = load_env_variables()
|
31 |
|
32 |
-
|
33 |
-
|
34 |
-
tokenizer = AutoTokenizer.from_pretrained('nvidia/NV-Embed-v1', token = hf_token , trust_remote_code=True)
|
35 |
-
model = AutoModel.from_pretrained('nvidia/NV-Embed-v1' , token = hf_token , trust_remote_code=True).to(device)
|
36 |
|
37 |
-
## add chroma vector store
|
38 |
|
39 |
-
##
|
40 |
|
41 |
-
|
42 |
api_key=yi_token,
|
43 |
base_url=API_BASE
|
44 |
)
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
49 |
# print(completion)
|
50 |
|
51 |
def respond(
|
|
|
18 |
from dotenv import load_dotenv
|
19 |
import re
|
20 |
from utils import load_env_variables
|
21 |
+
import chromadb
|
22 |
+
from chromadb import Documents, EmbeddingFunction, Embeddings
|
23 |
+
from chromadb.config import Settings
|
24 |
+
from chromadb import HttpClient
|
25 |
+
from langchain_community.document_loaders import UnstructuredFileLoader
|
26 |
+
from utils import load_env_variables , parse_and_route
|
27 |
|
28 |
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'
|
29 |
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
|
|
|
34 |
|
35 |
hf_token, yi_token = load_env_variables()
|
36 |
|
37 |
+
def clear_cuda_cache():
|
38 |
+
torch.cuda.empty_cache()
|
|
|
|
|
39 |
|
|
|
40 |
|
41 |
+
## 01ai Yi-large Clience
|
42 |
|
43 |
+
client = OpenAI(
|
44 |
api_key=yi_token,
|
45 |
base_url=API_BASE
|
46 |
)
|
47 |
+
|
48 |
+
|
49 |
+
## use instruct embeddings
|
50 |
+
|
51 |
+
# Load the tokenizer and model
|
52 |
+
|
53 |
+
class EmbeddingGenerator:
|
54 |
+
def __init__(self, model_name: str, token: str, intention_client):
|
55 |
+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
56 |
+
self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=token, trust_remote_code=True)
|
57 |
+
self.model = AutoModel.from_pretrained(model_name, token=token, trust_remote_code=True).to(self.device)
|
58 |
+
self.intention_client = intention_client
|
59 |
+
|
60 |
+
def clear_cuda_cache(self):
|
61 |
+
torch.cuda.empty_cache()
|
62 |
+
|
63 |
+
def compute_embeddings(self, input_text: str):
|
64 |
+
# Get the intention
|
65 |
+
intention_completion = self.intention_client.chat.completions.create(
|
66 |
+
model="yi-large",
|
67 |
+
messages=[
|
68 |
+
{"role": "system", "content": intention_prompt},
|
69 |
+
{"role": "user", "content": input_text}
|
70 |
+
]
|
71 |
+
)
|
72 |
+
intention_output = intention_completion.choices[0].message['content']
|
73 |
+
|
74 |
+
# Parse and route the intention
|
75 |
+
parsed_task = parse_and_route(intention_output)
|
76 |
+
selected_task = list(parsed_task.keys())[0]
|
77 |
+
|
78 |
+
# Construct the prompt
|
79 |
+
try:
|
80 |
+
task_description = tasks[selected_task]
|
81 |
+
except KeyError:
|
82 |
+
print(f"Selected task not found: {selected_task}")
|
83 |
+
return f"Error: Task '{selected_task}' not found. Please select a valid task."
|
84 |
+
|
85 |
+
query_prefix = f"Instruct: {task_description}\nQuery: "
|
86 |
+
queries = [input_text]
|
87 |
+
|
88 |
+
# Get the embeddings
|
89 |
+
with torch.no_grad():
|
90 |
+
inputs = self.tokenizer(queries, return_tensors='pt', padding=True, truncation=True, max_length=4096).to(self.device)
|
91 |
+
outputs = self.model(**inputs)
|
92 |
+
query_embeddings = outputs.last_hidden_state.mean(dim=1)
|
93 |
+
|
94 |
+
# Normalize embeddings
|
95 |
+
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
|
96 |
+
embeddings_list = query_embeddings.detach().cpu().numpy().tolist()
|
97 |
+
self.clear_cuda_cache()
|
98 |
+
return embeddings_list
|
99 |
+
|
100 |
+
|
101 |
+
class MyEmbeddingFunction(EmbeddingFunction):
|
102 |
+
def __init__(self, embedding_generator: EmbeddingGenerator):
|
103 |
+
self.embedding_generator = embedding_generator
|
104 |
+
|
105 |
+
def __call__(self, input: Documents) -> Embeddings:
|
106 |
+
embeddings = [self.embedding_generator.compute_embeddings(doc) for doc in input]
|
107 |
+
embeddings = [item for sublist in embeddings for item in sublist]
|
108 |
+
return embeddings
|
109 |
+
|
110 |
+
## add chroma vector store
|
111 |
+
class DocumentLoader:
|
112 |
+
def __init__(self, file_path: str, mode: str = "elements"):
|
113 |
+
self.file_path = file_path
|
114 |
+
self.mode = mode
|
115 |
+
|
116 |
+
def load_documents(self):
|
117 |
+
loader = UnstructuredFileLoader(self.file_path, mode=self.mode)
|
118 |
+
docs = loader.load()
|
119 |
+
return [doc.page_content for doc in docs]
|
120 |
+
|
121 |
+
class ChromaManager:
|
122 |
+
def __init__(self, collection_name: str, embedding_function: MyEmbeddingFunction):
|
123 |
+
self.client = HttpClient(settings=Settings(allow_reset=True))
|
124 |
+
self.client.reset() # resets the database
|
125 |
+
self.collection = self.client.create_collection(collection_name)
|
126 |
+
self.embedding_function = embedding_function
|
127 |
+
|
128 |
+
def add_documents(self, documents: list):
|
129 |
+
for doc in documents:
|
130 |
+
self.collection.add(ids=[str(uuid.uuid1())], documents=[doc], embeddings=self.embedding_function([doc]))
|
131 |
+
|
132 |
+
def query(self, query_text: str):
|
133 |
+
db = Chroma(client=self.client, collection_name=self.collection.name, embedding_function=self.embedding_function)
|
134 |
+
result_docs = db.similarity_search(query_text)
|
135 |
+
return result_docs
|
136 |
+
|
137 |
+
|
138 |
+
|
139 |
# print(completion)
|
140 |
|
141 |
def respond(
|
requirements.txt
CHANGED
@@ -4,4 +4,12 @@ sentence-transformers
|
|
4 |
torch==2.2.0
|
5 |
transformers
|
6 |
openai
|
7 |
-
python-dotenv
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
4 |
torch==2.2.0
|
5 |
transformers
|
6 |
openai
|
7 |
+
python-dotenv
|
8 |
+
chromadb
|
9 |
+
langchain-community
|
10 |
+
unstructured[all-docs]
|
11 |
+
libmagic
|
12 |
+
poppler
|
13 |
+
tesseract
|
14 |
+
libxml2
|
15 |
+
libxslt
|
utils.py
CHANGED
@@ -1,3 +1,5 @@
|
|
|
|
|
|
1 |
import re
|
2 |
from dotenv import load_dotenv
|
3 |
import re
|
@@ -14,18 +16,18 @@ def load_env_variables():
|
|
14 |
|
15 |
return hf_token, yi_token
|
16 |
|
17 |
-
def parse_and_route(example_output: str):
|
18 |
-
# Regex pattern to match the true task
|
19 |
-
pattern = r'"(\w+)":\s?true'
|
20 |
-
|
21 |
-
# Find the true task
|
22 |
-
match = re.search(pattern, example_output)
|
23 |
-
|
24 |
-
if match:
|
25 |
-
true_task = match.group(1)
|
26 |
-
if true_task in tasks:
|
27 |
-
return {true_task: tasks[true_task]}
|
28 |
-
else:
|
29 |
-
return {true_task: "Task description not found"}
|
30 |
-
else:
|
31 |
-
return "No true task found in the example output"
|
|
|
1 |
+
# utils.py
|
2 |
+
|
3 |
import re
|
4 |
from dotenv import load_dotenv
|
5 |
import re
|
|
|
16 |
|
17 |
return hf_token, yi_token
|
18 |
|
19 |
+
def parse_and_route(example_output: str):
|
20 |
+
# Regex pattern to match the true task
|
21 |
+
pattern = r'"(\w+)":\s?true'
|
22 |
+
|
23 |
+
# Find the true task
|
24 |
+
match = re.search(pattern, example_output)
|
25 |
+
|
26 |
+
if match:
|
27 |
+
true_task = match.group(1)
|
28 |
+
if true_task in tasks:
|
29 |
+
return {true_task: tasks[true_task]}
|
30 |
+
else:
|
31 |
+
return {true_task: "Task description not found"}
|
32 |
+
else:
|
33 |
+
return "No true task found in the example output"
|