Explorar el Código

Refactor: Improve the try catch logic for XinferenceEmbed (#9128)

### What problem does this PR solve?

Improve the try catch logic for XinferenceEmbed

### Type of change


- [x] Refactoring
tags/v0.20.0
Stephen Hu hace 3 meses
padre
commit
20b4d88098
No account linked to committer's email address
Se han modificado 1 ficheros con 4 adiciones y 2 borrados
  1. 4
    2
      rag/llm/embedding_model.py

+ 4
- 2
rag/llm/embedding_model.py Ver fichero

ress = [] ress = []
total_tokens = 0 total_tokens = 0
for i in range(0, len(texts), batch_size): for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i : i + batch_size], model=self.model_name)
res = None
try: try:
res = self.client.embeddings.create(input=texts[i : i + batch_size], model=self.model_name)
ress.extend([d.embedding for d in res.data]) ress.extend([d.embedding for d in res.data])
total_tokens += self.total_token_count(res) total_tokens += self.total_token_count(res)
except Exception as _e: except Exception as _e:
return np.array(ress), total_tokens return np.array(ress), total_tokens


def encode_queries(self, text): def encode_queries(self, text):
res = self.client.embeddings.create(input=[text], model=self.model_name)
res = None
try: try:
res = self.client.embeddings.create(input=[text], model=self.model_name)
return np.array(res.data[0].embedding), self.total_token_count(res) return np.array(res.data[0].embedding), self.total_token_count(res)
except Exception as _e: except Exception as _e:
log_exception(_e, res) log_exception(_e, res)

Cargando…
Cancelar
Guardar