Generate embedding Args: text: Text to embed model_name: Model name Returns: Embedding vector
(text: str, model_name: str = "text-embedding-3-small")
| 144 | |
| 145 | |
| 146 | def generate_embedding(text: str, model_name: str = "text-embedding-3-small") -> List[float]: |
| 147 | """ |
| 148 | Generate embedding |
| 149 | |
| 150 | Args: |
| 151 | text: Text to embed |
| 152 | model_name: Model name |
| 153 | |
| 154 | Returns: |
| 155 | Embedding vector |
| 156 | """ |
| 157 | # Check cache |
| 158 | cache_dir = Path(__file__).parent.parent.parent / "data" / "embeddings_cache" |
| 159 | cache_dir.mkdir(parents=True, exist_ok=True) |
| 160 | |
| 161 | text_hash = hashlib.sha256(text.encode()).hexdigest() |
| 162 | cache_file = cache_dir / f"{text_hash}.json" |
| 163 | |
| 164 | if cache_file.exists(): |
| 165 | with open(cache_file, 'r') as f: |
| 166 | cached = json.load(f) |
| 167 | if cached.get("model") == model_name: |
| 168 | return cached["embedding"] |
| 169 | |
| 170 | # Use OpenAI |
| 171 | if model_name.startswith("text-embedding") and HAS_OPENAI and client: |
| 172 | response = client.embeddings.create( |
| 173 | input=text, |
| 174 | model=model_name |
| 175 | ) |
| 176 | embedding = response.data[0].embedding |
| 177 | |
| 178 | # Cache |
| 179 | with open(cache_file, 'w') as f: |
| 180 | json.dump({ |
| 181 | "text_hash": text_hash, |
| 182 | "embedding": embedding, |
| 183 | "model": model_name, |
| 184 | "created_at": str(__import__("datetime").datetime.now()) |
| 185 | }, f) |
| 186 | |
| 187 | return embedding |
| 188 | |
| 189 | # Use local model |
| 190 | if HAS_LOCAL_MODEL and model and model_name in ["bge-large-zh", "m3e-base"]: |
| 191 | embedding = model.encode(text).tolist() |
| 192 | |
| 193 | # Cache |
| 194 | with open(cache_file, 'w') as f: |
| 195 | json.dump({ |
| 196 | "text_hash": text_hash, |
| 197 | "embedding": embedding, |
| 198 | "model": model_name, |
| 199 | "created_at": str(__import__("datetime").datetime.now()) |
| 200 | }, f) |
| 201 | |
| 202 | return embedding |
| 203 |