DEVOCEAN

class MultiQueriesRetriever:

	def __init__(self, base_retriever, model, model_name):
		self.base_retriever = base_retriever
		self.model = model
		self.model_name = model_name
	
	def _get_top_n_unique(self,topks, n):
		sorted_topks = sorted(topks, key=lambda x: x.score, reverse=True)
		unique_nodes = []
		seen_ids = set()
		for node_with_score in sorted_topks:
			if node_with_score.id_ not in seen_ids:
			unique_nodes.append(node_with_score)
			seen_ids.add(node_with_score.id_)
			if len(unique_nodes) == n:
				break
		return unique_nodes
	
	def _gen_queries(self, question) -> List[str]:
	
		prompt = f"""주어진 문장에 대해 동일한 내용을 갖는 유사한 한글 문장 5개를 생성합니다. json 형태로 생성해 주세요. json 내용이외의 내용은 생성하지 마세요.
		문장: {question}
		[
			{{ "문장": """
		
		res = self.model.chat.completions.create(
		model=self.model_name,
		messages=[
			{"role": "system", "content": "당신은 한글 생성 모델입니다. 요청한 내용에 맞게 문장을 생성해 주세요. "},
			{
			"role": "user",
			"content": prompt
			}
		],
		temperature=0.,
		max_tokens=512,
		top_p=1
		)
		
		questions=[question]
		try:
			responses=json.loads(res.choices[0].message.content)
			for i,element in enumerate(responses):
				questions.append(element['문장'])
		except:
			print(f'{res=} not proper response from llm')
		return questions
		
		def get_topk(self,question,n):
			questions=self._gen_queries(question)
			topks=[]
			for i,q in enumerate(questions):
				topks+=self.base_retriever.as_retriever(similarity_top_k=n).retrieve(q)
			topks=self._get_top_n_unique(topks,n)
			
			return topks