AI-B commited on
Commit
c4ddfdf
1 Parent(s): 9b90d69

Update app.py

Browse files

Changes Made:
Adjusted the __init__ method of MyEmbeddingFunction to accept model_name, token, and intention_client.
Modified the embedding_function initialization to match the new __init__ method signature.
This should resolve the TypeError you encountered. Let me know if you need any further assistance!

Files changed (1) hide show
  1. app.py +9 -13
app.py CHANGED
@@ -1,4 +1,4 @@
1
- # app.py
2
  import spaces
3
  from torch.nn import DataParallel
4
  from torch import Tensor
@@ -17,12 +17,8 @@ import gradio as gr
17
  import torch
18
  import torch.nn.functional as F
19
  from dotenv import load_dotenv
20
- from utils import load_env_variables, parse_and_route , escape_special_characters
21
- from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name , metadata_prompt
22
- # import time
23
- # import httpx
24
-
25
-
26
 
27
  load_dotenv()
28
 
@@ -110,7 +106,7 @@ class EmbeddingGenerator:
110
  matches = pattern.findall(metadata_output)
111
  metadata = {key: value for key, value in matches}
112
  return metadata
113
-
114
  class MyEmbeddingFunction(EmbeddingFunction):
115
  def __init__(self, model_name: str, token: str, intention_client):
116
  self.model_name = model_name
@@ -140,7 +136,7 @@ def initialize_chroma(collection_name: str, embedding_function: MyEmbeddingFunct
140
 
141
  def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunction):
142
  for doc in documents:
143
- embeddings, metadata = embedding_function.embedding_generator.compute_embeddings(doc)
144
  for embedding, meta in zip(embeddings, metadata):
145
  chroma_collection.add(
146
  ids=[str(uuid.uuid1())],
@@ -150,7 +146,7 @@ def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunc
150
  )
151
 
152
  def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):
153
- query_embeddings, query_metadata = embedding_function.embedding_generator.compute_embeddings(query_text)
154
  result_docs = chroma_collection.query(
155
  query_texts=[query_text],
156
  n_results=2
@@ -160,7 +156,7 @@ def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):
160
  # Initialize clients
161
  intention_client = OpenAI(api_key=yi_token, base_url=API_BASE)
162
  embedding_generator = EmbeddingGenerator(model_name=model_name, token=hf_token, intention_client=intention_client)
163
- embedding_function = MyEmbeddingFunction(embedding_generator=embedding_generator)
164
  chroma_db = initialize_chroma(collection_name="Tonic-instruct", embedding_function=embedding_function)
165
 
166
  def respond(
@@ -199,7 +195,7 @@ def upload_documents(files):
199
  return "Documents uploaded and processed successfully!"
200
 
201
  def query_documents(query):
202
- results = query_chroma(query)
203
  return "\n\n".join([result.content for result in results])
204
 
205
  with gr.Blocks() as demo:
@@ -226,4 +222,4 @@ with gr.Blocks() as demo:
226
 
227
  if __name__ == "__main__":
228
  # os.system("chroma run --host localhost --port 8000 &")
229
- demo.launch()
 
1
+ # app.py
2
  import spaces
3
  from torch.nn import DataParallel
4
  from torch import Tensor
 
17
  import torch
18
  import torch.nn.functional as F
19
  from dotenv import load_dotenv
20
+ from utils import load_env_variables, parse_and_route, escape_special_characters
21
+ from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name, metadata_prompt
 
 
 
 
22
 
23
  load_dotenv()
24
 
 
106
  matches = pattern.findall(metadata_output)
107
  metadata = {key: value for key, value in matches}
108
  return metadata
109
+
110
  class MyEmbeddingFunction(EmbeddingFunction):
111
  def __init__(self, model_name: str, token: str, intention_client):
112
  self.model_name = model_name
 
136
 
137
  def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunction):
138
  for doc in documents:
139
+ embeddings, metadata = embedding_function.create_embedding_generator().compute_embeddings(doc)
140
  for embedding, meta in zip(embeddings, metadata):
141
  chroma_collection.add(
142
  ids=[str(uuid.uuid1())],
 
146
  )
147
 
148
  def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):
149
+ query_embeddings, query_metadata = embedding_function.create_embedding_generator().compute_embeddings(query_text)
150
  result_docs = chroma_collection.query(
151
  query_texts=[query_text],
152
  n_results=2
 
156
  # Initialize clients
157
  intention_client = OpenAI(api_key=yi_token, base_url=API_BASE)
158
  embedding_generator = EmbeddingGenerator(model_name=model_name, token=hf_token, intention_client=intention_client)
159
+ embedding_function = MyEmbeddingFunction(model_name=model_name, token=hf_token, intention_client=intention_client)
160
  chroma_db = initialize_chroma(collection_name="Tonic-instruct", embedding_function=embedding_function)
161
 
162
  def respond(
 
195
  return "Documents uploaded and processed successfully!"
196
 
197
  def query_documents(query):
198
+ results = query_chroma(query, embedding_function)
199
  return "\n\n".join([result.content for result in results])
200
 
201
  with gr.Blocks() as demo:
 
222
 
223
  if __name__ == "__main__":
224
  # os.system("chroma run --host localhost --port 8000 &")
225
+ demo.launch()