return fq and model names mixed together from fuzzy match

This commit is contained in:
Paul Gauthier 2024-07-14 18:26:05 +01:00
parent 29afc85957
commit 50353acb1a

View file

@ -654,11 +654,7 @@ def sanity_check_model(io, model):
if possible_matches:
io.tool_output("Did you mean one of these?")
for match in possible_matches:
fq, m = match
if fq == m:
io.tool_output(f"- {m}")
else:
io.tool_output(f"- {m} ({fq})")
io.tool_output(f"- {model}")
if show:
io.tool_output(f"For more info, see: {urls.model_warnings}\n")
@ -667,7 +663,7 @@ def sanity_check_model(io, model):
def fuzzy_match_models(name):
name = name.lower()
chat_models = []
chat_models = set()
for model, attrs in litellm.model_cost.items():
model = model.lower()
if attrs.get("mode") != "chat":
@ -679,8 +675,10 @@ def fuzzy_match_models(name):
else:
fq_model = provider + model
chat_models.append((fq_model, model))
chat_models.add(fq_model)
chat_models.add(model)
chat_models = sorted(chat_models)
# exactly matching model
# matching_models = [
# (fq,m) for fq,m in chat_models
@ -690,19 +688,14 @@ def fuzzy_match_models(name):
# return matching_models
# Check for model names containing the name
matching_models = [(fq, m) for fq, m in chat_models if name in fq]
matching_models = [m for m in chat_models if name in m]
if matching_models:
return matching_models
# Check for slight misspellings
models = [m for fq, m in chat_models]
models = list(chat_models)
matching_models = difflib.get_close_matches(name, models, n=3, cutoff=0.8)
if matching_models:
return list(zip(matching_models, matching_models))
fq_models = [fq for fq, m in chat_models]
matching_models = difflib.get_close_matches(name, fq_models, n=3, cutoff=0.8)
return list(zip(matching_models, matching_models))
return sorted(matching_models)
def print_matching_models(io, search):
@ -710,8 +703,7 @@ def print_matching_models(io, search):
if matches:
io.tool_output(f'Models which match "{search}":')
for model in matches:
fq, m = model
io.tool_output(f"- {fq}")
io.tool_output(f"- {model}")
else:
io.tool_output(f'No models match "{search}".')