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

Refactor code formatting in lightrag_api_openai_compatible_demo.py

Browse Source
This commit is contained in:
zhenjie.ye
2024-10-26 16:12:10 +08:00
parent 981be9e569
commit 883d6b7cc7
3 changed files with 28 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lightrag/lightrag.py
View File

@@ -85,9 +85,7 @@ class LightRAG:
# LLM # LLM
llm_model_func: callable = gpt_4o_mini_complete # hf_model_complete# llm_model_func: callable = gpt_4o_mini_complete # hf_model_complete#
llm_model_name: str = ( llm_model_name: str = "meta-llama/Llama-3.2-1B-Instruct" #'meta-llama/Llama-3.2-1B'#'google/gemma-2-2b-it'
"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_token_size: int = 32768
llm_model_max_async: int = 16 llm_model_max_async: int = 16

4
lightrag/llm.py
View File

@@ -286,7 +286,9 @@ async def hf_model_if_cache(
output = hf_model.generate( output = hf_model.generate(
**input_ids, max_new_tokens=512, num_return_sequences=1, early_stopping=True **input_ids, max_new_tokens=512, num_return_sequences=1, early_stopping=True
) )
response_text = hf_tokenizer.decode(output[0][len(inputs["input_ids"][0]):], skip_special_tokens=True) response_text = hf_tokenizer.decode(
output[0][len(inputs["input_ids"][0]) :], skip_special_tokens=True
)
if hashing_kv is not None: if hashing_kv is not None:
await hashing_kv.upsert({args_hash: {"return": response_text, "model": model}}) await hashing_kv.upsert({args_hash: {"return": response_text, "model": model}})
return response_text return response_text

27
setup.py
View File

@@ -1,6 +1,7 @@
import setuptools import setuptools
from pathlib import Path from pathlib import Path
# Reading the long description from README.md # Reading the long description from README.md
def read_long_description(): def read_long_description():
try: try:
@@ -8,6 +9,7 @@ def read_long_description():
except FileNotFoundError: except FileNotFoundError:
return "A description of LightRAG is currently unavailable." return "A description of LightRAG is currently unavailable."
# Retrieving metadata from __init__.py # Retrieving metadata from __init__.py
def retrieve_metadata(): def retrieve_metadata():
vars2find = ["__author__", "__version__", "__url__"] vars2find = ["__author__", "__version__", "__url__"]
@@ -17,7 +19,12 @@ def retrieve_metadata():
for line in f.readlines(): for line in f.readlines():
for v in vars2find: for v in vars2find:
if line.startswith(v): if line.startswith(v):
line = line.replace(" ", "").replace('"', "").replace("'", "").strip() line = (
line.replace(" ", "")
.replace('"', "")
.replace("'", "")
.strip()
)
vars2readme[v] = line.split("=")[1] vars2readme[v] = line.split("=")[1]
except FileNotFoundError: except FileNotFoundError:
raise FileNotFoundError("Metadata file './lightrag/__init__.py' not found.") raise FileNotFoundError("Metadata file './lightrag/__init__.py' not found.")
@@ -25,10 +32,13 @@ def retrieve_metadata():
# Checking if all required variables are found # Checking if all required variables are found
missing_vars = [v for v in vars2find if v not in vars2readme] missing_vars = [v for v in vars2find if v not in vars2readme]
if missing_vars: if missing_vars:
raise ValueError(f"Missing required metadata variables in __init__.py: {missing_vars}") raise ValueError(
f"Missing required metadata variables in __init__.py: {missing_vars}"
)
return vars2readme return vars2readme
# Reading dependencies from requirements.txt # Reading dependencies from requirements.txt
def read_requirements(): def read_requirements():
deps = [] deps = []
@@ -36,9 +46,12 @@ def read_requirements():
with open("./requirements.txt") as f: with open("./requirements.txt") as f:
deps = [line.strip() for line in f if line.strip()] deps = [line.strip() for line in f if line.strip()]
except FileNotFoundError: except FileNotFoundError:
print("Warning: 'requirements.txt' not found. No dependencies will be installed.") print(
"Warning: 'requirements.txt' not found. No dependencies will be installed."
)
return deps return deps
metadata = retrieve_metadata() metadata = retrieve_metadata()
long_description = read_long_description() long_description = read_long_description()
requirements = read_requirements() requirements = read_requirements()
@@ -51,7 +64,9 @@ setuptools.setup(
description="LightRAG: Simple and Fast Retrieval-Augmented Generation", description="LightRAG: Simple and Fast Retrieval-Augmented Generation",
long_description=long_description, long_description=long_description,
long_description_content_type="text/markdown", long_description_content_type="text/markdown",
packages=setuptools.find_packages(exclude=("tests*", "docs*")), # Automatically find packages packages=setuptools.find_packages(
exclude=("tests*", "docs*")
), # Automatically find packages
classifiers=[ classifiers=[
"Development Status :: 4 - Beta", "Development Status :: 4 - Beta",
"Programming Language :: Python :: 3", "Programming Language :: Python :: 3",
@@ -66,6 +81,8 @@ setuptools.setup(
project_urls={ # Additional project metadata project_urls={ # Additional project metadata
"Documentation": metadata.get("__url__", ""), "Documentation": metadata.get("__url__", ""),
"Source": metadata.get("__url__", ""), "Source": metadata.get("__url__", ""),
"Tracker": f"{metadata.get('__url__', '')}/issues" if metadata.get("__url__") else "" "Tracker": f"{metadata.get('__url__', '')}/issues"
if metadata.get("__url__")
else "",
}, },
) )