Ley_Fill7 commited on
Commit
a9c90e6
·
1 Parent(s): 9ed0ff7

Added a way to resolve callback manager in embed model

Browse files
Files changed (1) hide show
  1. app.py +7 -1
app.py CHANGED
@@ -2,6 +2,7 @@
2
  from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage
3
  from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings, NVIDIARerank
4
  from llama_index.core.indices.query.query_transform import HyDEQueryTransform
 
5
  from llama_index.core.query_engine import TransformQueryEngine
6
  from langchain_core.documents import Document as LangDocument
7
  from llama_index.core import Document as LlamaDocument
@@ -29,13 +30,18 @@ embed_model = NVIDIAEmbeddings(
29
  truncate="NONE"
30
  )
31
 
 
 
 
 
 
32
  reranker = NVIDIARerank(
33
  model="nvidia/nv-rerankqa-mistral-4b-v3",
34
  api_key=nvidia_api_key,
35
  )
36
 
37
  # Set the NVIDIA models globally
38
- Settings.embed_model = embed_model
39
  Settings.llm = client
40
 
41
  # Parse the local PDF document
 
2
  from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage
3
  from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings, NVIDIARerank
4
  from llama_index.core.indices.query.query_transform import HyDEQueryTransform
5
+ from llama_index.core.embeddings.utils import resolve_embed_model
6
  from llama_index.core.query_engine import TransformQueryEngine
7
  from langchain_core.documents import Document as LangDocument
8
  from llama_index.core import Document as LlamaDocument
 
30
  truncate="NONE"
31
  )
32
 
33
+ def custom_resolve_embed_model(embed_model):
34
+ if isinstance(embed_model, NVIDIAEmbeddings):
35
+ return embed_model
36
+ return resolve_embed_model(embed_model)
37
+
38
  reranker = NVIDIARerank(
39
  model="nvidia/nv-rerankqa-mistral-4b-v3",
40
  api_key=nvidia_api_key,
41
  )
42
 
43
  # Set the NVIDIA models globally
44
+ Settings.embed_model = custom_resolve_embed_model(embed_model)
45
  Settings.llm = client
46
 
47
  # Parse the local PDF document