LocalAI/scripts/model_gallery_info.py
Ettore Di Giacinto e8128a339a
chore(scripts): handle summarization errors (#4271)
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
2024-11-26 14:51:55 +01:00

116 lines
3.5 KiB
Python

## This script simply help pull off some info from the HF api
## to speed up addition of new models to the gallery.
## It accepts as input a repo_id and returns part of the YAML data
## Use it as:
## OPENAI_BASE_URL="<api_url>" OPENAI_MODEL="" python .github/add_model.py mradermacher/HaloMaidRP-v1.33-15B-L3-i1-GGUF
## Example:
# local-ai run hermes-2-theta-llama-3-8b
# OPENAI_BASE_URL="http://192.168.xx.xx:8080" OPENAI_MODEL="hermes-2-theta-llama-3-8b" python scripts/model_gallery_info.py mradermacher/HaloMaidRP-v1.33-15B-L3-i1-GGUF
import sys
import os
from openai import OpenAI
from huggingface_hub import HfFileSystem, get_paths_info
templated_yaml = """
- !!merge <<: *llama3
name: "{model_name}"
urls:
- https://huggingface.co/{repo_id}
description: |
{description}
overrides:
parameters:
model: {file_name}
files:
- filename: {file_name}
sha256: {checksum}
uri: huggingface://{repo_id}/{file_name}
"""
client = OpenAI()
model = os.environ.get("OPENAI_MODEL", "hermes-2-theta-llama-3-8b")
def summarize(text: str) -> str:
chat_completion = client.chat.completions.create(
messages=[
{
"role": "user",
"content": "You are a bot which extracts the description of the LLM model from the following text. Return ONLY the description of the model, and nothing else.\n" + text,
},
],
model=model,
)
return chat_completion.choices[0].message.content
def format_description(description):
return '\n '.join(description.split('\n'))
# Example usage
if __name__ == "__main__":
# Get repoid from argv[0]
repo_id = sys.argv[1]
token = "" # Replace with your Hugging Face token if needed
fs = HfFileSystem()
all_files = fs.ls(repo_id, detail=False)
print(all_files)
# Find a file that has Q4_K in the name
file_path = None
file_name = None
readmeFile = None
for file in all_files:
print(f"File found: {file}")
if "readme" in file.lower():
readmeFile = file
print(f"Found README file: {readmeFile}")
if "q4_k_m" in file.lower():
file_path = file
if file_path is None:
print("No file with Q4_K_M found, using the first file in the list.")
exit(1)
# Extract file from full path (is the last element)
if file_path is not None:
file_name = file_path.split("/")[-1]
model_name = repo_id.split("/")[-1]
checksum = None
for file in get_paths_info(repo_id, [file_name], repo_type='model'):
try:
checksum = file.lfs.sha256
break
except Exception as e:
print(f'Error from Hugging Face Hub: {str(e)}', file=sys.stderr)
sys.exit(2)
print(checksum)
print(file_name)
print(file_path)
summarized_readme = ""
if readmeFile:
# If there is a README file, read it
readme = fs.read_text(readmeFile)
try:
summarized_readme = summarize(readme)
except Exception as e:
print(f"Error summarizing the README: {str(e)}", file=sys.stderr)
summarized_readme = format_description(summarized_readme)
print("Model correctly processed")
## Append to the result YAML file
with open("result.yaml", "a") as f:
f.write(templated_yaml.format(model_name=model_name.lower().replace("-GGUF","").replace("-gguf",""), repo_id=repo_id, description=summarized_readme, file_name=file_name, checksum=checksum, file_path=file_path))