Tonic commited on
Commit
eb88ab8
1 Parent(s): 2ca300b

add intention mapper to embedding flow , refactor into classes

Browse files
Files changed (3) hide show
  1. app.py +102 -12
  2. requirements.txt +9 -1
  3. 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
- ## use instruct embeddings
33
- # Load the tokenizer and model
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
- ## Make intention Mapper
40
 
41
- intention_client = OpenAI(
42
  api_key=yi_token,
43
  base_url=API_BASE
44
  )
45
- intention_completion = intention_client.chat.completions.create(
46
- model="yi-large",
47
- messages=[{"role": "system", "content": intention_prompt},{"role": "user", "content": inputext}]
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"