|
|
|
@ -4,18 +4,19 @@ https://arxiv.org/abs/2212.10496
|
|
|
|
|
"""
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
|
|
from typing import List
|
|
|
|
|
from typing import Dict, List
|
|
|
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
|
from pydantic import BaseModel, Extra
|
|
|
|
|
|
|
|
|
|
from langchain.chains.base import Chain
|
|
|
|
|
from langchain.chains.hyde.prompts import PROMPT_MAP
|
|
|
|
|
from langchain.chains.llm import LLMChain
|
|
|
|
|
from langchain.embeddings.base import Embeddings
|
|
|
|
|
from langchain.embeddings.hyde.prompts import PROMPT_MAP
|
|
|
|
|
from langchain.llms.base import BaseLLM
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class HypotheticalDocumentEmbedder(Embeddings, BaseModel):
|
|
|
|
|
class HypotheticalDocumentEmbedder(Chain, Embeddings, BaseModel):
|
|
|
|
|
"""Generate hypothetical document for query, and then embed that.
|
|
|
|
|
|
|
|
|
|
Based on https://arxiv.org/abs/2212.10496
|
|
|
|
@ -30,10 +31,24 @@ class HypotheticalDocumentEmbedder(Embeddings, BaseModel):
|
|
|
|
|
extra = Extra.forbid
|
|
|
|
|
arbitrary_types_allowed = True
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
def input_keys(self) -> List[str]:
|
|
|
|
|
"""Input keys for Hyde's LLM chain."""
|
|
|
|
|
return self.llm_chain.input_keys
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
def output_keys(self) -> List[str]:
|
|
|
|
|
"""Output keys for Hyde's LLM chain."""
|
|
|
|
|
return self.llm_chain.output_keys
|
|
|
|
|
|
|
|
|
|
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
|
|
|
|
"""Call the base embeddings."""
|
|
|
|
|
return self.base_embeddings.embed_documents(texts)
|
|
|
|
|
|
|
|
|
|
def combine_embeddings(self, embeddings: List[List[float]]) -> List[float]:
|
|
|
|
|
"""Combine embeddings into final embeddings."""
|
|
|
|
|
return list(np.array(embeddings).mean(axis=0))
|
|
|
|
|
|
|
|
|
|
def embed_query(self, text: str) -> List[float]:
|
|
|
|
|
"""Generate a hypothetical document and embedded it."""
|
|
|
|
|
var_name = self.llm_chain.input_keys[0]
|
|
|
|
@ -42,9 +57,9 @@ class HypotheticalDocumentEmbedder(Embeddings, BaseModel):
|
|
|
|
|
embeddings = self.embed_documents(documents)
|
|
|
|
|
return self.combine_embeddings(embeddings)
|
|
|
|
|
|
|
|
|
|
def combine_embeddings(self, embeddings: List[List[float]]) -> List[float]:
|
|
|
|
|
"""Combine embeddings into final embeddings."""
|
|
|
|
|
return list(np.array(embeddings).mean(axis=0))
|
|
|
|
|
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
|
|
|
|
|
"""Call the internal llm chain."""
|
|
|
|
|
return self.llm_chain._call(inputs)
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
def from_llm(
|