File size: 2,445 Bytes
b5deaf1 ad04a72 b5deaf1 e83b975 ad04a72 af61c79 ad04a72 83747e9 ad04a72 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
"""Module for OpenAI model and embeddings."""
from typing import List
from langchain.embeddings.base import Embeddings
from sentence_transformers import SentenceTransformer
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
class GPTModel(AzureChatOpenAI):
"""
GPTModel class that extends AzureChatOpenAI.
This class initializes a GPT model with specific deployment settings and a callback function.
Attributes:
callback (function): The callback function to be used with the model.
Methods:
__init__(callback):
Initializes the GPTModel with the specified callback function.
"""
def __init__(self):
super().__init__(
deployment_name="gpt-4o",
streaming=True, temperature=0)
class GPTEmbeddings(AzureOpenAIEmbeddings):
"""
GPTEmbeddings class that extends AzureOpenAIEmbeddings.
This class is designed to handle embeddings using GPT model provided by Azure OpenAI services.
Attributes:
Inherits all attributes from AzureOpenAIEmbeddings.
Methods:
Inherits all methods from AzureOpenAIEmbeddings.
"""
class EmbeddingsModel(Embeddings):
"""
A model for generating embeddings using SentenceTransformer.
Attributes:
model (SentenceTransformer): The SentenceTransformer model used for generating embeddings.
"""
def __init__(self, model_name: str):
"""
Initializes the Chroma model with the specified model name.
Args:
model_name (str): The name of the model to be used for embedding.
"""
self.model = SentenceTransformer(model_name)
def embed_documents(self, documents: List[str]) -> List[List[float]]:
"""
Embed a list of documents into a list of vectors.
Args:
documents (List[str]): A list of documents to be embedded.
Returns:
List[List[float]]: A list of vectors representing the embedded documents.
"""
return self.model.encode(documents).tolist()
def embed_query(self, query: str) -> List[float]:
"""
Embed a query string into a list of floats using the model's encoding.
Args:
query (str): The query string to be embedded.
Returns:
List[float]: The embedded representation of the query as a list of floats.
"""
return self.model.encode([query]).tolist()[0]
|