mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-02-04 01:25:07 +08:00
Don't release full image (#10654)
### What problem does this PR solve? Introduced gpu profile in .env Added Dockerfile_tei fix datrie Removed LIGHTEN flag ### Type of change - [x] Documentation Update - [x] Refactoring
This commit is contained in:
@ -1,109 +0,0 @@
|
||||
#
|
||||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
from jina import Deployment
|
||||
from docarray import BaseDoc
|
||||
from jina import Executor, requests
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
|
||||
import argparse
|
||||
import torch
|
||||
|
||||
|
||||
class Prompt(BaseDoc):
|
||||
message: list[dict]
|
||||
gen_conf: dict
|
||||
|
||||
|
||||
class Generation(BaseDoc):
|
||||
text: str
|
||||
|
||||
|
||||
tokenizer = None
|
||||
model_name = ""
|
||||
|
||||
|
||||
class TokenStreamingExecutor(Executor):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name, device_map="auto", torch_dtype="auto"
|
||||
)
|
||||
|
||||
@requests(on="/chat")
|
||||
async def generate(self, doc: Prompt, **kwargs) -> Generation:
|
||||
text = tokenizer.apply_chat_template(
|
||||
doc.message,
|
||||
tokenize=False,
|
||||
)
|
||||
inputs = tokenizer([text], return_tensors="pt")
|
||||
generation_config = GenerationConfig(
|
||||
**doc.gen_conf,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
pad_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
generated_ids = self.model.generate(
|
||||
inputs.input_ids, generation_config=generation_config
|
||||
)
|
||||
generated_ids = [
|
||||
output_ids[len(input_ids) :]
|
||||
for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
|
||||
]
|
||||
|
||||
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
||||
yield Generation(text=response)
|
||||
|
||||
@requests(on="/stream")
|
||||
async def task(self, doc: Prompt, **kwargs) -> Generation:
|
||||
text = tokenizer.apply_chat_template(
|
||||
doc.message,
|
||||
tokenize=False,
|
||||
)
|
||||
input = tokenizer([text], return_tensors="pt")
|
||||
input_len = input["input_ids"].shape[1]
|
||||
max_new_tokens = 512
|
||||
if "max_new_tokens" in doc.gen_conf:
|
||||
max_new_tokens = doc.gen_conf.pop("max_new_tokens")
|
||||
generation_config = GenerationConfig(
|
||||
**doc.gen_conf,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
pad_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
for _ in range(max_new_tokens):
|
||||
output = self.model.generate(
|
||||
**input, max_new_tokens=1, generation_config=generation_config
|
||||
)
|
||||
if output[0][-1] == tokenizer.eos_token_id:
|
||||
break
|
||||
yield Generation(
|
||||
text=tokenizer.decode(output[0][input_len:], skip_special_tokens=True)
|
||||
)
|
||||
input = {
|
||||
"input_ids": output,
|
||||
"attention_mask": torch.ones(1, len(output[0])),
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model_name", type=str, help="Model name or path")
|
||||
parser.add_argument("--port", default=12345, type=int, help="Jina serving port")
|
||||
args = parser.parse_args()
|
||||
model_name = args.model_name
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
|
||||
with Deployment(
|
||||
uses=TokenStreamingExecutor, port=args.port, protocol="grpc"
|
||||
) as dep:
|
||||
dep.block()
|
||||
@ -29,6 +29,7 @@ from api.db.services.pipeline_operation_log_service import PipelineOperationLogS
|
||||
from api.utils.api_utils import timeout
|
||||
from api.utils.base64_image import image2id
|
||||
from api.utils.log_utils import init_root_logger, get_project_base_directory
|
||||
from api.utils.configs import show_configs
|
||||
from graphrag.general.index import run_graphrag_for_kb
|
||||
from graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache
|
||||
from rag.flow.pipeline import Pipeline
|
||||
@ -475,7 +476,7 @@ async def embedding(docs, mdl, parser_config=None, callback=None):
|
||||
tk_count = 0
|
||||
if len(tts) == len(cnts):
|
||||
vts, c = await trio.to_thread.run_sync(lambda: mdl.encode(tts[0: 1]))
|
||||
tts = np.concatenate([vts for _ in range(len(tts))], axis=0)
|
||||
tts = np.concatenate([vts[0] for _ in range(len(tts))], axis=0)
|
||||
tk_count += c
|
||||
|
||||
@timeout(60)
|
||||
@ -1061,7 +1062,10 @@ async def main():
|
||||
/____/
|
||||
""")
|
||||
logging.info(f'RAGFlow version: {get_ragflow_version()}')
|
||||
show_configs()
|
||||
settings.init_settings()
|
||||
from api.settings import EMBEDDING_CFG
|
||||
logging.info(f'api.settings.EMBEDDING_CFG: {EMBEDDING_CFG}')
|
||||
print_rag_settings()
|
||||
if sys.platform != "win32":
|
||||
signal.signal(signal.SIGUSR1, start_tracemalloc_and_snapshot)
|
||||
|
||||
Reference in New Issue
Block a user