gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add HF Support

Browse Source
This commit is contained in:
TianyuFan0504
2024-10-14 19:41:07 +08:00
parent 665f50c8fe
commit 741953c34b
3 changed files with 148 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
lightrag/lightrag.py
View File

@@ -5,7 +5,7 @@ from datetime import datetime
from functools import partial
from typing import Type, cast
from .llm import gpt_4o_complete, gpt_4o_mini_complete, openai_embedding
from .llm import gpt_4o_complete, gpt_4o_mini_complete, openai_embedding,hf_model,hf_embedding
from .operate import (
chunking_by_token_size,
extract_entities,
@@ -77,12 +77,13 @@ class LightRAG:
)
# text embedding
embedding_func: EmbeddingFunc = field(default_factory=lambda: openai_embedding)
embedding_func: EmbeddingFunc = field(default_factory=lambda: hf_embedding)#openai_embedding
embedding_batch_num: int = 32
embedding_func_max_async: int = 16
# LLM
llm_model_func: callable = gpt_4o_mini_complete
llm_model_func: callable = hf_model#gpt_4o_mini_complete
llm_model_name: str = 'meta-llama/Llama-3.2-1B-Instruct'#'meta-llama/Llama-3.2-1B'#'google/gemma-2-2b-it'
llm_model_max_token_size: int = 32768
llm_model_max_async: int = 16