OpenSourceAppConfig.py 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. from typing import Optional
  2. from chromadb.utils import embedding_functions
  3. from embedchain.helper_classes.json_serializable import register_deserializable
  4. from .BaseAppConfig import BaseAppConfig
  5. @register_deserializable
  6. class OpenSourceAppConfig(BaseAppConfig):
  7. """
  8. Config to initialize an embedchain custom `OpenSourceApp` instance, with extra config options.
  9. """
  10. def __init__(
  11. self,
  12. log_level=None,
  13. host=None,
  14. port=None,
  15. id=None,
  16. collection_name=None,
  17. collect_metrics: Optional[bool] = None,
  18. model=None,
  19. ):
  20. """
  21. :param log_level: Optional. (String) Debug level
  22. ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'].
  23. :param id: Optional. ID of the app. Document metadata will have this id.
  24. :param collection_name: Optional. Collection name for the database.
  25. :param host: Optional. Hostname for the database server.
  26. :param port: Optional. Port for the database server.
  27. :param collect_metrics: Defaults to True. Send anonymous telemetry to improve embedchain.
  28. :param model: Optional. GPT4ALL uses the model to instantiate the class.
  29. So unlike `App`, it has to be provided before querying.
  30. """
  31. self.model = model or "orca-mini-3b.ggmlv3.q4_0.bin"
  32. super().__init__(
  33. log_level=log_level,
  34. embedding_fn=OpenSourceAppConfig.default_embedding_function(),
  35. host=host,
  36. port=port,
  37. id=id,
  38. collection_name=collection_name,
  39. collect_metrics=collect_metrics,
  40. )
  41. @staticmethod
  42. def default_embedding_function():
  43. """
  44. Sets embedding function to default (`all-MiniLM-L6-v2`).
  45. :returns: The default embedding function
  46. """
  47. try:
  48. return embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
  49. except ValueError as e:
  50. print(e)
  51. raise ModuleNotFoundError(
  52. "The open source app requires extra dependencies. Install with `pip install embedchain[opensource]`"
  53. ) from None