Spaces:
Running
Running
| import asyncio | |
| import os | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dotenv import load_dotenv | |
| from sentence_transformers import ( | |
| CrossEncoder, | |
| ) # SentenceTransformer -> model for embeddings, CrossEncoder -> re-ranker | |
| from torch import Tensor | |
| from google import genai | |
| from google.genai import types | |
| from app.core.chunks import Chunk | |
| from app.settings import settings, BASE_DIR, GeminiEmbeddingSettings, logger | |
| import aiofiles.os | |
| load_dotenv() | |
| class Reranker: | |
| def __init__(self, model: str = "cross-encoder/ms-marco-MiniLM-L6-v2"): | |
| self.device: str = settings.device | |
| self.model_name: str = model | |
| self.model: CrossEncoder = CrossEncoder(model, device=self.device) | |
| async def rank(self, query: str, chunks: list[Chunk]) -> list[dict[str, int]]: | |
| return await asyncio.to_thread(self.model.rank, query, [await chunk.get_raw_text() for chunk in chunks]) | |
| class GeminiLLM: | |
| def __init__(self, model="gemini-2.0-flash"): | |
| self.client = genai.Client(api_key=settings.api_key) | |
| self.model = model | |
| async def get_response(self, prompt: str, stream: bool = True, use_default_config: bool = False) -> str: | |
| path_to_prompt = os.path.join(BASE_DIR, "models_io", "prompt.txt") | |
| async with aiofiles.open(path_to_prompt, "w", encoding="utf-8", errors="replace") as f: | |
| await f.write(prompt) | |
| response = self.client.models.generate_content( | |
| model=self.model, | |
| contents=prompt, | |
| config=( | |
| types.GenerateContentConfig(**settings.gemini_generation.model_dump()) | |
| if use_default_config | |
| else None | |
| ), | |
| ) | |
| return response.text | |
| async def get_streaming_response(self, prompt: str, use_default_config: bool = False): | |
| loop = asyncio.get_event_loop() | |
| start = loop.time() | |
| path_to_prompt = os.path.join(BASE_DIR, "models_io", "prompt.txt") | |
| async with aiofiles.open(path_to_prompt, "w", encoding="utf-8", errors="replace") as f: | |
| await f.write(prompt) | |
| await logger.info(f"Time of saving prompt to document - {loop.time() - start}") | |
| response = self.client.models.generate_content_stream( | |
| model=self.model, | |
| contents=prompt, | |
| config=( | |
| types.GenerateContentConfig(**settings.gemini_generation.model_dump()) | |
| if use_default_config | |
| else None | |
| ), | |
| ) | |
| for chunk in response: | |
| yield chunk | |
| class GeminiEmbed: | |
| def __init__(self, model="text-embedding-004"): | |
| self.client = genai.Client(api_key=settings.api_key) | |
| self.model = model | |
| self.settings = GeminiEmbeddingSettings() | |
| self.max_workers = 5 | |
| self.embed_executor = ThreadPoolExecutor(max_workers=self.max_workers) | |
| def _embed_batch_sync(self, batch: list[str], idx: int) -> dict: | |
| response = self.client.models.embed_content( | |
| model=self.model, | |
| contents=batch, | |
| config=types.EmbedContentConfig( | |
| **settings.gemini_embedding.model_dump() | |
| ), | |
| ).embeddings | |
| return {"idx": idx, "embeddings": response} | |
| async def encode(self, text: str | list[str]) -> list[Tensor]: | |
| if isinstance(text, str): | |
| text = [text] | |
| groups: list[dict] = [] | |
| max_batch_size = 100 # can not be changed due to google restrictions | |
| batches: list[list[str]] = [ | |
| text[i : i + max_batch_size] | |
| for i in range(0, len(text), max_batch_size) | |
| ] | |
| loop = asyncio.get_running_loop() | |
| tasks = [ | |
| loop.run_in_executor( | |
| self.embed_executor, | |
| self._embed_batch_sync, | |
| batch, | |
| idx | |
| ) | |
| for idx, batch in enumerate(batches) | |
| ] | |
| groups = await asyncio.gather(*tasks) | |
| groups.sort(key=lambda x: x["idx"]) | |
| result: list[float] = [] | |
| for group in groups: | |
| for vec in group["embeddings"]: | |
| result.append(vec.values) | |
| return result | |
| async def get_vector_dimensionality(self) -> int | None: | |
| return getattr(self.settings, "output_dimensionality") | |
| class Wrapper: | |
| def __init__(self, model: str = "gemini-2.0-flash"): | |
| self.model = model | |
| self.client = genai.Client(api_key=settings.api_key) | |
| async def wrap(self, prompt: str) -> str: | |
| def wrapper(prompt): | |
| response = self.client.models.generate_content( | |
| model=self.model, | |
| contents=prompt, | |
| config=types.GenerateContentConfig(**settings.gemini_wrapper.model_dump()) | |
| ) | |
| return response.text | |
| return await asyncio.to_thread(wrapper, prompt) | |