change embedding api

Author: shahules786

nbs/embedding/base.ipynb

@@ -46,7 +46,7 @@
     "from tqdm import tqdm\n",
     "\n",
     "from ragas_annotator.prompt.base import Prompt\n",
-    "from ragas_annotator.embedding.base import RagasEmbedding\n",
+    "from ragas_annotator.embedding.base import BaseEmbedding\n",
     "from ragas_annotator.metric import MetricResult\n",
     "from ragas_annotator.llm import RagasLLM\n",
     "from ragas_annotator.project.core import Project\n",
@@ -119,7 +119,7 @@
     "        # Run all tasks concurrently and return results\n",
     "        return await asyncio.gather(*async_tasks)\n",
     "    \n",
-    "    def train(self,project:Project, experiment_names: t.List[str], model:NotionModel, embedding_model: RagasEmbedding,method: t.Dict[str, t.Any]):\n",
+    "    def train(self,project:Project, experiment_names: t.List[str], model:NotionModel, embedding_model: BaseEmbedding,method: t.Dict[str, t.Any]):\n",
     "        \n",
     "        assert isinstance(self.prompt, Prompt)\n",
     "        self.prompt = DynamicFewShotPrompt.from_prompt(self.prompt,embedding_model)\n",

nbs/metric/base.ipynb

@@ -61,24 +61,28 @@
     "    \n",
     "    def format(self, **kwargs) -> str:\n",
     "        \"\"\"Format the prompt with the provided variables.\"\"\"\n",
-    "        # Start with examples if we have them\n",
+    "\n",
     "        prompt_parts = []\n",
-    "        \n",
-    "                \n",
-    "        # Add instruction with variables filled in\n",
     "        prompt_parts.append(self.instruction.format(**kwargs))\n",
+    "        prompt_parts.append(self._format_examples())\n",
+    "\n",
+    "        # Combine all parts\n",
+    "        return \"\\n\\n\".join(prompt_parts)\n",
+    "    \n",
+    "    def _format_examples(self) -> str:\n",
     "        \n",
     "        # Add examples in a simple format\n",
+    "        examples = []\n",
     "        if self.examples:\n",
-    "            prompt_parts.append(\"Examples:\")\n",
+    "            examples.append(\"Examples:\")\n",
     "            for i, (inputs, output) in enumerate(self.examples, 1):\n",
     "                example_input = \"\\n\".join([f\"{k}: {v}\" for k, v in inputs.items()])\n",
     "                example_output = \"\\n\".join([f\"{k}: {v}\" for k, v in output.items()])\n",
     "                \n",
-    "                prompt_parts.append(f\"Example {i}:\\nInput:\\n{example_input}\\nOutput:\\n{example_output}\")\n",
-    "\n",
-    "        # Combine all parts\n",
-    "        return \"\\n\\n\".join(prompt_parts)\n",
+    "                examples.append(f\"Example {i}:\\nInput:\\n{example_input}\\nOutput:\\n{example_output}\")\n",
+    "        \n",
+    "        return \"\\n\\n\".join(examples) if examples else \"\"\n",
+    "        \n",
     "    \n",
     "    def add_example(self, inputs: t.Dict, output: t.Dict) -> None:\n",
     "        \"\"\"\n",
@@ -105,7 +109,7 @@
     "    \n",
     "    def __str__(self) -> str:\n",
     "        \"\"\"String representation showing the instruction.\"\"\"\n",
-    "        return f\"Prompt(instruction='{self.instruction}')\""
+    "        return f\"Prompt(instruction='{self.instruction}',\\n examples={self.examples})\""
    ]
   },
   {
@@ -169,6 +173,45 @@
     "\n",
     "print(prompt.format(response=\"You can get a full refund if you miss your flight.\", expected_answer=\"Refunds depend on ticket type; only refundable tickets qualify for full refunds.\"))"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Prompt(instruction='Evaluate if given answer {response} is same as expected answer {expected_answer}',\n",
+      " examples=Examples:\n",
+      "\n",
+      "Example 1:\n",
+      "Input:\n",
+      "response: You can get a full refund if you miss your flight.\n",
+      "expected_answer: Refunds depend on ticket type; only refundable tickets qualify for full refunds.\n",
+      "Output:\n",
+      "score: fail\n",
+      "\n",
+      "Example 2:\n",
+      "Input:\n",
+      "response: Each passenger gets 1 free checked bag up to 23kg.\n",
+      "expected_answer: Each passenger gets 1 free checked bag up to 23kg.\n",
+      "Output:\n",
+      "score: pass)\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(str(prompt))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

nbs/prompt/base.ipynb

@@ -29,7 +29,7 @@
     "from abc import ABC, abstractmethod\n",
     "\n",
     "from ragas_annotator.prompt.base import Prompt\n",
-    "from ragas_annotator.embedding.base import RagasEmbedding\n",
+    "from ragas_annotator.embedding import BaseEmbedding\n",
     "\n",
     "class ExampleStore(ABC):\n",
     "    @abstractmethod\n",
@@ -205,7 +205,7 @@
     "    def from_prompt(\n",
     "        cls,\n",
     "        prompt: Prompt,\n",
-    "        embedding_model: RagasEmbedding,\n",
+    "        embedding_model: BaseEmbedding,\n",
     "        num_examples: int = 3\n",
     "    ) -> \"DynamicFewShotPrompt\":\n",
     "        \"\"\"Create a DynamicFewShotPrompt from a Prompt object.\"\"\"\n",
@@ -251,10 +251,11 @@
    ],
    "source": [
     "#| eval: false\n",
+    "from ragas_annotator.embedding import ragas_embedding\n",
     "from ragas_annotator.prompt import Prompt\n",
     "from openai import OpenAI\n",
     "\n",
-    "embedding = RagasEmbedding(client=OpenAI(),model=\"text-embedding-3-small\")\n",
+    "embedding = ragas_embedding(provider=\"openai\", client=OpenAI(),model=\"text-embedding-3-small\")\n",
     "\n",
     "# Create a basic prompt\n",
     "prompt = Prompt(\n",

nbs/prompt/dynamic_few_shot.ipynb

@@ -119,16 +119,30 @@
                                          'ragas_annotator.dataset.Dataset.pop': ('dataset.html#dataset.pop', 'ragas_annotator/dataset.py'),
                                          'ragas_annotator.dataset.Dataset.save': ( 'dataset.html#dataset.save',
                                                                                    'ragas_annotator/dataset.py')},
-            'ragas_annotator.embedding.base': { 'ragas_annotator.embedding.base.RagasEmbedding': ( 'embedding/base.html#ragasembedding',
-                                                                                                   'ragas_annotator/embedding/base.py'),
-                                                'ragas_annotator.embedding.base.RagasEmbedding.aembed_document': ( 'embedding/base.html#ragasembedding.aembed_document',
-                                                                                                                   'ragas_annotator/embedding/base.py'),
-                                                'ragas_annotator.embedding.base.RagasEmbedding.aembed_text': ( 'embedding/base.html#ragasembedding.aembed_text',
-                                                                                                               'ragas_annotator/embedding/base.py'),
-                                                'ragas_annotator.embedding.base.RagasEmbedding.embed_document': ( 'embedding/base.html#ragasembedding.embed_document',
+            'ragas_annotator.embedding.base': { 'ragas_annotator.embedding.base.BaseEmbedding': ( 'embedding/base.html#baseembedding',
+                                                                                                  'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.BaseEmbedding.aembed_document': ( 'embedding/base.html#baseembedding.aembed_document',
                                                                                                                   'ragas_annotator/embedding/base.py'),
-                                                'ragas_annotator.embedding.base.RagasEmbedding.embed_text': ( 'embedding/base.html#ragasembedding.embed_text',
-                                                                                                              'ragas_annotator/embedding/base.py')},
+                                                'ragas_annotator.embedding.base.BaseEmbedding.aembed_text': ( 'embedding/base.html#baseembedding.aembed_text',
+                                                                                                              'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.BaseEmbedding.embed_document': ( 'embedding/base.html#baseembedding.embed_document',
+                                                                                                                 'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.BaseEmbedding.embed_text': ( 'embedding/base.html#baseembedding.embed_text',
+                                                                                                             'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.OpenAIEmbeddings': ( 'embedding/base.html#openaiembeddings',
+                                                                                                     'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.OpenAIEmbeddings.__init__': ( 'embedding/base.html#openaiembeddings.__init__',
+                                                                                                              'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.OpenAIEmbeddings.aembed_document': ( 'embedding/base.html#openaiembeddings.aembed_document',
+                                                                                                                     'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.OpenAIEmbeddings.aembed_text': ( 'embedding/base.html#openaiembeddings.aembed_text',
+                                                                                                                 'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.OpenAIEmbeddings.embed_document': ( 'embedding/base.html#openaiembeddings.embed_document',
+                                                                                                                    'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.OpenAIEmbeddings.embed_text': ( 'embedding/base.html#openaiembeddings.embed_text',
+                                                                                                                'ragas_annotator/embedding/base.py'),
+                                                'ragas_annotator.embedding.base.ragas_embedding': ( 'embedding/base.html#ragas_embedding',
+                                                                                                    'ragas_annotator/embedding/base.py')},
             'ragas_annotator.exceptions': { 'ragas_annotator.exceptions.DuplicateError': ( 'utils/exceptions.html#duplicateerror',
                                                                                            'ragas_annotator/exceptions.py'),
                                             'ragas_annotator.exceptions.NotFoundError': ( 'utils/exceptions.html#notfounderror',
@@ -417,6 +431,8 @@
                                                                                               'ragas_annotator/prompt/base.py'),
                                              'ragas_annotator.prompt.base.Prompt.__str__': ( 'prompt/base.html#prompt.__str__',
                                                                                              'ragas_annotator/prompt/base.py'),
+                                             'ragas_annotator.prompt.base.Prompt._format_examples': ( 'prompt/base.html#prompt._format_examples',
+                                                                                                      'ragas_annotator/prompt/base.py'),
                                              'ragas_annotator.prompt.base.Prompt._validate_instruction': ( 'prompt/base.html#prompt._validate_instruction',
                                                                                                            'ragas_annotator/prompt/base.py'),
                                              'ragas_annotator.prompt.base.Prompt.add_example': ( 'prompt/base.html#prompt.add_example',

ragas_annotator/_modidx.py

@@ -1,3 +1,4 @@
-from ragas_annotator.embedding.base import RagasEmbedding
+from ragas_annotator.embedding.base import BaseEmbedding
+from ragas_annotator.embedding.base import ragas_embedding
 
-__all__ = ['RagasEmbedding']
+__all__ = ['ragas_embedding','BaseEmbedding']

ragas_annotator/embedding/__init__.py

@@ -1,32 +1,67 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/embedding/base.ipynb.
 
 # %% auto 0
-__all__ = ['RagasEmbedding']
+__all__ = ['BaseEmbedding', 'OpenAIEmbeddings', 'ragas_embedding']
 
 # %% ../../nbs/embedding/base.ipynb 2
 import typing as t
-from dataclasses import dataclass
+from abc import ABC, abstractmethod
 
+#TODO: Add support for other providers like HuggingFace, Cohere, etc.
+#TODO: handle async calls properly and ensure that the client supports async if needed.
 
-@dataclass
-class RagasEmbedding:
-    client: t.Any
-    model: str
+class BaseEmbedding(ABC):
+    @abstractmethod
+    def embed_text(self, text: str, **kwargs: t.Any) -> t.List[float]:
+        pass
 
-    def embed_text(self,text:str,**kwargs: t.Any) -> t.List[float]:
-        
+    @abstractmethod
+    async def aembed_text(self, text: str, **kwargs: t.Any) -> t.List[float]:
+        pass
+    
+    @abstractmethod
+    def embed_document(self, documents: t.List[str], **kwargs: t.Any) -> t.List[t.List[float]]:
+        pass
+    
+    @abstractmethod
+    async def aembed_document(self, documents: t.List[str], **kwargs: t.Any) -> t.List[t.List[float]]:
+        pass
+
+
+class OpenAIEmbeddings(BaseEmbedding):
+    def __init__(self, client: t.Any, model: str):
+        self.client = client
+        self.model = model
+    
+    def embed_text(self, text: str, **kwargs: t.Any) -> t.List[float]:
         return self.client.embeddings.create(input=text, model=self.model, **kwargs).data[0].embedding
-        
-    async def aembed_text(self,text:str,**kwargs: t.Any):
-        
-        await self.client.embeddings.create(input=text, model=self.model, **kwargs).data[0].embedding
 
+    async def aembed_text(self, text: str, **kwargs: t.Any) -> t.List[float]:
+        response = await self.client.embeddings.create(input=text, model=self.model, **kwargs)
+        return response.data[0].embedding
 
-    def embed_document(self,documents:t.List[str],**kwargs: t.Any) -> t.List[t.List[float]]:
+    def embed_document(self, documents: t.List[str], **kwargs: t.Any) -> t.List[t.List[float]]:
         embeddings = self.client.embeddings.create(input=documents, model=self.model, **kwargs)
         return [embedding.embedding for embedding in embeddings.data]
 
-    async def aembed_document(self,documents:t.List[str],**kwargs: t.Any)  -> t.List[t.List[float]]:
+    async def aembed_document(self, documents: t.List[str], **kwargs: t.Any) -> t.List[t.List[float]]:
         embeddings = await self.client.embeddings.create(input=documents, model=self.model, **kwargs)
         return [embedding.embedding for embedding in embeddings.data]
-
+    
+    
+def ragas_embedding(provider: str, model: str, client: t.Any) -> BaseEmbedding:
+    """
+    Factory function to create an embedding instance based on the provider.
+    
+    Args:
+        provider (str): The name of the embedding provider (e.g., "openai").
+        model (str): The model name to use for embeddings.
+        **kwargs: Additional arguments for the provider's client.
+    
+    Returns:
+        BaseEmbedding: An instance of the specified embedding provider.
+    """
+    if provider.lower() == "openai":
+        return OpenAIEmbeddings(client=client, model=model)
+    
+    raise ValueError(f"Unsupported provider: {provider}")

ragas_annotator/embedding/base.py

@@ -15,7 +15,7 @@
 from tqdm import tqdm
 
 from ..prompt.base import Prompt
-from ..embedding.base import RagasEmbedding
+from ..embedding.base import BaseEmbedding
 from . import MetricResult
 from ..llm import RagasLLM
 from ..project.core import Project
@@ -88,7 +88,7 @@ async def abatch_score(self, inputs: t.List[t.Dict[str, t.Any]], reasoning: bool
         # Run all tasks concurrently and return results
         return await asyncio.gather(*async_tasks)
 
-    def train(self,project:Project, experiment_names: t.List[str], model:NotionModel, embedding_model: RagasEmbedding,method: t.Dict[str, t.Any]):
+    def train(self,project:Project, experiment_names: t.List[str], model:NotionModel, embedding_model: BaseEmbedding,method: t.Dict[str, t.Any]):
 
         assert isinstance(self.prompt, Prompt)
         self.prompt = DynamicFewShotPrompt.from_prompt(self.prompt,embedding_model)

ragas_annotator/metric/base.py

@@ -41,24 +41,28 @@ def _validate_instruction(self):
 
     def format(self, **kwargs) -> str:
         """Format the prompt with the provided variables."""
-        # Start with examples if we have them
+
         prompt_parts = []
-        
-                
-        # Add instruction with variables filled in
         prompt_parts.append(self.instruction.format(**kwargs))
+        prompt_parts.append(self._format_examples())
+
+        # Combine all parts
+        return "\n\n".join(prompt_parts)
+    
+    def _format_examples(self) -> str:
 
         # Add examples in a simple format
+        examples = []
         if self.examples:
-            prompt_parts.append("Examples:")
+            examples.append("Examples:")
             for i, (inputs, output) in enumerate(self.examples, 1):
                 example_input = "\n".join([f"{k}: {v}" for k, v in inputs.items()])
                 example_output = "\n".join([f"{k}: {v}" for k, v in output.items()])
 
-                prompt_parts.append(f"Example {i}:\nInput:\n{example_input}\nOutput:\n{example_output}")
-
-        # Combine all parts
-        return "\n\n".join(prompt_parts)
+                examples.append(f"Example {i}:\nInput:\n{example_input}\nOutput:\n{example_output}")
+        
+        return "\n\n".join(examples) if examples else ""
+        
 
     def add_example(self, inputs: t.Dict, output: t.Dict) -> None:
         """
@@ -85,4 +89,4 @@ def add_example(self, inputs: t.Dict, output: t.Dict) -> None:
 
     def __str__(self) -> str:
         """String representation showing the instruction."""
-        return f"Prompt(instruction='{self.instruction}')"
+        return f"Prompt(instruction='{self.instruction}',\n examples={self.examples})"

ragas_annotator/prompt/base.py

@@ -9,7 +9,7 @@
 from abc import ABC, abstractmethod
 
 from .base import Prompt
-from ..embedding.base import RagasEmbedding
+from ..embedding import BaseEmbedding
 
 class ExampleStore(ABC):
     @abstractmethod
@@ -185,7 +185,7 @@ def add_example(self, inputs: t.Dict, output: t.Dict) -> None:
     def from_prompt(
         cls,
         prompt: Prompt,
-        embedding_model: RagasEmbedding,
+        embedding_model: BaseEmbedding,
         num_examples: int = 3
     ) -> "DynamicFewShotPrompt":
         """Create a DynamicFewShotPrompt from a Prompt object."""