pdf file is downloadable, embeddings are local

end-to-end-use-cases/Contextual-Chunking-RAG/README.md

@@ -11,8 +11,7 @@
 **Getting started**
 In this cookbook, we’ll use DeepInfra for Llama inference services, so be sure to obtain an API key from https://deepinfra.com/.
 You'll also need a LlamaParse API key to parse PDF files, which can be obtained from https://www.llamaindex.ai/.
-Additionally, an OpenAI API key is required for embeddings.
+Additionally, we will use the "jinaai/jina-embeddings-v2-base-en" model from HuggingFace to generate text embeddings locally.
 Before getting started, update the <code>config.py</code> file as following:
-    "DEEPINFRA_API_KEY"="<your_api_key>"
-    "OPENAI_API_KEY"="<your_api_key>"
+    "DEEPINFRA_API_KEY"="<your_api_key>"    
     "LLAMAPARSE_API_KEY"="<your_api_key>"

end-to-end-use-cases/Contextual-Chunking-RAG/config.py

@@ -1,5 +1,2 @@
-OPENAI_API_KEY=""
 LLAMAPARSE_API_KEY=""
 DEEPINFRA_API_KEY=""
-
-

end-to-end-use-cases/Contextual-Chunking-RAG/embedding.py

@@ -0,0 +1,43 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+from llama_index.core.base.embeddings.base import BaseEmbedding
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load tokenizer and model
+model_id = "jinaai/jina-embeddings-v2-base-en" #"jinaai/jina-embeddings-v3"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device)
+
+# Define function to generate embeddings
+def get_embedding(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)    
+    return outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy() #.to(torch.float32)
+
+
+class LocalJinaEmbedding(BaseEmbedding):
+    def __init__(self):
+        super().__init__()
+
+    def _get_text_embedding(self, text):
+        return get_embedding(text).tolist()  # Ensure compatibility with LlamaIndex
+
+    def _get_query_embedding(self, query):
+        return get_embedding(query).tolist()
+    
+    async def _aget_query_embedding(self, query: str) -> list:
+        return get_embedding(query).tolist()
+
+
+
+def test(): #this did not produce reasonable results for some reason
+    #!pip install llama-index-embeddings-huggingface
+    from llama_index.embeddings.huggingface import HuggingFaceEmbedding 
+    embed_model = HuggingFaceEmbedding(model_name=model_id)
+
+
+if __name__=="__main__":
+	emb = get_embedding("hi there")
+	print(emb.shape)