pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refine resume parts and fix bugs in retrival using sql (#66)

Browse Source
This commit is contained in:
KevinHuSh
2024-02-19 19:22:17 +08:00
committed by GitHub
parent 452020d33a
commit a8294f2168
29 changed files with 302 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
rag/llm/__init__.py
View File

@ -21,7 +21,7 @@ from .cv_model import *
EmbeddingModel = {
"Infiniflow": HuEmbedding,
"OpenAI": OpenAIEmbed,
"通义千问": QWenEmbed,
"通义千问": HuEmbedding, #QWenEmbed,
}

12
rag/llm/chat_model.py
View File

@ -32,7 +32,7 @@ class GptTurbo(Base):
self.model_name = model_name
def chat(self, system, history, gen_conf):
history.insert(0, {"role": "system", "content": system})
if system: history.insert(0, {"role": "system", "content": system})
res = self.client.chat.completions.create(
model=self.model_name,
messages=history,
@ -49,11 +49,12 @@ class QWenChat(Base):
def chat(self, system, history, gen_conf):
from http import HTTPStatus
history.insert(0, {"role": "system", "content": system})
if system: history.insert(0, {"role": "system", "content": system})
response = Generation.call(
self.model_name,
messages=history,
result_format='message'
result_format='message',
**gen_conf
)
if response.status_code == HTTPStatus.OK:
return response.output.choices[0]['message']['content'], response.usage.output_tokens
@ -68,10 +69,11 @@ class ZhipuChat(Base):
def chat(self, system, history, gen_conf):
from http import HTTPStatus
history.insert(0, {"role": "system", "content": system})
if system: history.insert(0, {"role": "system", "content": system})
response = self.client.chat.completions.create(
self.model_name,
messages=history
messages=history,
**gen_conf
)
if response.status_code == HTTPStatus.OK:
return response.output.choices[0]['message']['content'], response.usage.completion_tokens

6
rag/llm/embedding_model.py
View File

@ -100,11 +100,11 @@ class QWenEmbed(Base):
input=texts[i:i+batch_size],
text_type="document"
)
embds = [[]] * len(resp["output"]["embeddings"])
embds = [[] for _ in range(len(resp["output"]["embeddings"]))]
for e in resp["output"]["embeddings"]:
embds[e["text_index"]] = e["embedding"]
res.extend(embds)
token_count += resp["usage"]["input_tokens"]
token_count += resp["usage"]["total_tokens"]
return np.array(res), token_count
def encode_queries(self, text):
@ -113,7 +113,7 @@ class QWenEmbed(Base):
input=text[:2048],
text_type="query"
)
return np.array(resp["output"]["embeddings"][0]["embedding"]), resp["usage"]["input_tokens"]
return np.array(resp["output"]["embeddings"][0]["embedding"]), resp["usage"]["total_tokens"]
from zhipuai import ZhipuAI