Spaces:
Running
Running
File size: 1,356 Bytes
559d75f |
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 |
# retriever.py
from typing import Any, Dict, List
import torch
from langchain.tools.retriever import create_retriever_tool
class MultiModalRetriever:
"""
Enhanced retrieval system that integrates text, image, and code snippet search.
"""
def __init__(self, text_retriever: Any, clip_model: Any, clip_processor: Any) -> None:
self.text_retriever = text_retriever
self.clip_model = clip_model
self.clip_processor = clip_processor
self.code_retriever = create_retriever_tool([], "Code Retriever", "Retriever for code snippets")
def retrieve(self, query: str, domain: str) -> Dict[str, List]:
return {
"text": self._retrieve_text(query),
"images": self._retrieve_images(query),
"code": self._retrieve_code(query)
}
def _retrieve_text(self, query: str) -> List[Any]:
return self.text_retriever.invoke(query)
def _retrieve_images(self, query: str) -> List[str]:
inputs = self.clip_processor(text=query, return_tensors="pt")
with torch.no_grad():
_ = self.clip_model.get_text_features(**inputs)
# Placeholder for image retrieval results
return ["image_result_1.png", "image_result_2.png"]
def _retrieve_code(self, query: str) -> List[str]:
return self.code_retriever.invoke(query)
|