MCPcopy Create free account
hub / github.com/NanGePlus/LightRAGTest / hf_model_if_cache

Function hf_model_if_cache

LightRAG/lightrag/llm.py:235–297  ·  view source on GitHub ↗
(
    model, prompt, system_prompt=None, history_messages=[], **kwargs
)

Source from the content-addressed store, hash-verified

233
234
235async def hf_model_if_cache(
236 model, prompt, system_prompt=None, history_messages=[], **kwargs
237) -> str:
238 model_name = model
239 hf_model, hf_tokenizer = initialize_hf_model(model_name)
240 hashing_kv: BaseKVStorage = kwargs.pop("hashing_kv", None)
241 messages = []
242 if system_prompt:
243 messages.append({"role": "system", "content": system_prompt})
244 messages.extend(history_messages)
245 messages.append({"role": "user", "content": prompt})
246
247 if hashing_kv is not None:
248 args_hash = compute_args_hash(model, messages)
249 if_cache_return = await hashing_kv.get_by_id(args_hash)
250 if if_cache_return is not None:
251 return if_cache_return["return"]
252 input_prompt = ""
253 try:
254 input_prompt = hf_tokenizer.apply_chat_template(
255 messages, tokenize=False, add_generation_prompt=True
256 )
257 except Exception:
258 try:
259 ori_message = copy.deepcopy(messages)
260 if messages[0]["role"] == "system":
261 messages[1]["content"] = (
262 "<system>"
263 + messages[0]["content"]
264 + "</system>\n"
265 + messages[1]["content"]
266 )
267 messages = messages[1:]
268 input_prompt = hf_tokenizer.apply_chat_template(
269 messages, tokenize=False, add_generation_prompt=True
270 )
271 except Exception:
272 len_message = len(ori_message)
273 for msgid in range(len_message):
274 input_prompt = (
275 input_prompt
276 + "<"
277 + ori_message[msgid]["role"]
278 + ">"
279 + ori_message[msgid]["content"]
280 + "</"
281 + ori_message[msgid]["role"]
282 + ">\n"
283 )
284
285 input_ids = hf_tokenizer(
286 input_prompt, return_tensors="pt", padding=True, truncation=True
287 ).to("cuda")
288 inputs = {k: v.to(hf_model.device) for k, v in input_ids.items()}
289 output = hf_model.generate(
290 **input_ids, max_new_tokens=512, num_return_sequences=1, early_stopping=True
291 )
292 response_text = hf_tokenizer.decode(

Callers 1

hf_model_completeFunction · 0.85

Calls 4

initialize_hf_modelFunction · 0.85
compute_args_hashFunction · 0.85
get_by_idMethod · 0.45
upsertMethod · 0.45

Tested by

no test coverage detected