fix formating

Author: jjmachan

src/ragas/executor.py

@@ -3,7 +3,6 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 
-import numpy as np
 from tqdm.auto import tqdm
 
 

src/ragas/llms/base.py

@@ -36,11 +36,10 @@ def is_multiple_completion_supported(llm: BaseLanguageModel) -> bool:
 
 @dataclass
 class BaseRagasLLM(ABC):
-    
     def get_temperature(self, n: int) -> float:
         """Return the temperature to use for completion based on n."""
         return 0.3 if n > 1 else 1e-8
-    
+
     @abstractmethod
     def generate_text(
         self,

src/ragas/testset/docstore.py

@@ -1,6 +1,6 @@
 import heapq
-import typing as t
 import logging
+import typing as t
 import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -11,8 +11,8 @@
 import numpy.typing as npt
 from langchain.text_splitter import TextSplitter
 from langchain_core.documents import Document as LCDocument
-from llama_index.readers.schema import Document as LlamaindexDocument
 from langchain_core.pydantic_v1 import Field
+from llama_index.readers.schema import Document as LlamaindexDocument
 
 from ragas.async_utils import run_async_tasks
 from ragas.embeddings.base import BaseRagasEmbeddings, embedding_factory

src/ragas/testset/evolutions.py

@@ -1,17 +1,17 @@
-from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
 import logging
 import typing as t
-from fsspec.exceptions import asyncio
-from random import choice, choices
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from random import choice
 
+from fsspec.exceptions import asyncio
 from langchain.prompts import ChatPromptTemplate
 from numpy.random import default_rng
 
 from ragas.llms import BaseRagasLLM
 from ragas.llms.json_load import load_as_json
 from ragas.llms.prompt import PromptValue
-from ragas.testset.docstore import Document, DocumentStore, Direction
+from ragas.testset.docstore import Direction, Document, DocumentStore, Node
 from ragas.testset.prompts import (
     FILTER_QUESTION,
     MULTICONTEXT_QUESTION,
@@ -20,7 +20,6 @@
     TABLE_QA,
     demonstrations,
 )
-from ragas.testset.docstore import Node
 
 rng = default_rng()
 logger = logging.getLogger(__name__)

src/ragas/testset/generator.py

@@ -3,13 +3,12 @@
 
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
-from ragas.llms import BaseRagasLLM, LangchainLLMWrapper
-from ragas.embeddings import BaseRagasEmbeddings
-from ragas.testset.docstore import DocumentStore, Document
-from ragas.testset.evolutions import SimpleEvolution
-
 from llama_index.readers.schema import Document as LlamaindexDocument
 
+from ragas.embeddings import BaseRagasEmbeddings
+from ragas.llms import BaseRagasLLM, LangchainLLMWrapper
+from ragas.testset.docstore import Document, DocumentStore, InMemoryDocumentStore
+
 
 @dataclass
 class TestsetGenerator:
@@ -24,15 +23,31 @@ def with_openai(
         generator_llm: str = "gpt-3.5-turbo",
         critic_llm: str = "gpt-4",
         embeddings: str = "text-embedding-ada-002",
+        docstore: t.Optional[DocumentStore] = None,
+        chunk_size: int = 512,
     ) -> "TestsetGenerator":
         generator_llm_model = LangchainLLMWrapper(ChatOpenAI(model=generator_llm))
         critic_llm_model = LangchainLLMWrapper(ChatOpenAI(model=critic_llm))
         embeddings_model = OpenAIEmbeddings(model=embeddings)
-        return cls(
-            generator_llm=generator_llm_model,
-            critic_llm=critic_llm_model,
-            embeddings=embeddings_model,
-        )
+        if docstore is None:
+            from langchain.text_splitter import TokenTextSplitter
+
+            splitter = TokenTextSplitter(chunk_size=chunk_size, chunk_overlap=0)
+            docstore = InMemoryDocumentStore(splitter)
+            return cls(
+                generator_llm=generator_llm_model,
+                critic_llm=critic_llm_model,
+                # TODO: remove type ignore after fixing embeddigns
+                embeddings=embeddings_model,  # type: ignore
+                docstore=docstore,
+            )
+        else:
+            return cls(
+                generator_llm=generator_llm_model,
+                critic_llm=critic_llm_model,
+                embeddings=embeddings_model,  # type: ignore
+                docstore=docstore,
+            )
 
     def generate_with_llamaindex_docs(self, documents: t.Sequence[LlamaindexDocument]):
         # chunk documents and add to docstore

src/ragas/testset/testset_generator.py

@@ -544,4 +544,4 @@ def generate(
                     count += 1
                     pbar.update(count)
 
-        return TestDataset(test_data=samples)
+        return TestDataset(test_data=samples)

tests/e2e/test_adaptation.py

@@ -1,4 +1,3 @@
-
 from ragas import adapt
 from ragas.metrics import context_recall
 

tests/unit/testset_generator/test_docstore.py

@@ -5,7 +5,7 @@
 import pytest
 from langchain_core.embeddings import Embeddings
 
-from ragas.testset.docstore import InMemoryDocumentStore, Node, Direction
+from ragas.testset.docstore import Direction, InMemoryDocumentStore, Node
 
 
 def test_adjacent_nodes():