Feature: Add openai_api_model attribute to Doctran models (#7868)

- Description: Added the ability to define the open AI model.
- Issue: Currently the Doctran instance uses gpt-4 by default, this does
not work if the user has no access to gpt -4.
  - rlancemartin, @eyurtsev, @baskaryan

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
Aiden Le 2023-07-20 09:27:56 -05:00 committed by GitHub
parent d8c40253c3
commit 1d7414a371
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 28 additions and 5 deletions

View File

@ -53,11 +53,15 @@ class DoctranPropertyExtractor(BaseDocumentTransformer):
self,
properties: List[dict],
openai_api_key: Optional[str] = None,
openai_api_model: Optional[str] = None,
) -> None:
self.properties = properties
self.openai_api_key = openai_api_key or get_from_env(
"openai_api_key", "OPENAI_API_KEY"
)
self.openai_api_model = openai_api_model or get_from_env(
"openai_api_model", "OPENAI_API_MODEL"
)
def transform_documents(
self, documents: Sequence[Document], **kwargs: Any
@ -71,7 +75,9 @@ class DoctranPropertyExtractor(BaseDocumentTransformer):
try:
from doctran import Doctran, ExtractProperty
doctran = Doctran(openai_api_key=self.openai_api_key)
doctran = Doctran(
openai_api_key=self.openai_api_key, openai_model=self.openai_api_model
)
except ImportError:
raise ImportError(
"Install doctran to use this parser. (pip install doctran)"

View File

@ -21,10 +21,17 @@ class DoctranQATransformer(BaseDocumentTransformer):
transformed_document = await qa_transformer.atransform_documents(documents)
"""
def __init__(self, openai_api_key: Optional[str] = None) -> None:
def __init__(
self,
openai_api_key: Optional[str] = None,
openai_api_model: Optional[str] = None,
) -> None:
self.openai_api_key = openai_api_key or get_from_env(
"openai_api_key", "OPENAI_API_KEY"
)
self.openai_api_model = openai_api_model or get_from_env(
"openai_api_model", "OPENAI_API_MODEL"
)
def transform_documents(
self, documents: Sequence[Document], **kwargs: Any
@ -38,7 +45,9 @@ class DoctranQATransformer(BaseDocumentTransformer):
try:
from doctran import Doctran
doctran = Doctran(openai_api_key=self.openai_api_key)
doctran = Doctran(
openai_api_key=self.openai_api_key, openai_model=self.openai_api_model
)
except ImportError:
raise ImportError(
"Install doctran to use this parser. (pip install doctran)"

View File

@ -23,11 +23,17 @@ class DoctranTextTranslator(BaseDocumentTransformer):
"""
def __init__(
self, openai_api_key: Optional[str] = None, language: str = "english"
self,
openai_api_key: Optional[str] = None,
language: str = "english",
openai_api_model: Optional[str] = None,
) -> None:
self.openai_api_key = openai_api_key or get_from_env(
"openai_api_key", "OPENAI_API_KEY"
)
self.openai_api_model = openai_api_model or get_from_env(
"openai_api_model", "OPENAI_API_MODEL"
)
self.language = language
def transform_documents(
@ -42,7 +48,9 @@ class DoctranTextTranslator(BaseDocumentTransformer):
try:
from doctran import Doctran
doctran = Doctran(openai_api_key=self.openai_api_key)
doctran = Doctran(
openai_api_key=self.openai_api_key, openai_model=self.openai_api_model
)
except ImportError:
raise ImportError(
"Install doctran to use this parser. (pip install doctran)"