Spaces:
Sleeping
Sleeping
File size: 6,600 Bytes
67beed8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
import json
from typing import Optional, List
import httpx
from common import LlmPredictParams, LlmParams, LlmApi
class DeepInfraApi(LlmApi):
"""
Класс для работы с API vllm.
"""
def __init__(self, params: LlmParams):
super.params = params
async def get_models(self) -> List[str]:
"""
Выполняет GET-запрос к API для получения списка доступных моделей.
Возвращает:
list[str]: Список идентификаторов моделей.
Если произошла ошибка или данные недоступны, возвращается пустой список.
Исключения:
Все ошибки HTTP-запросов логируются в консоль, но не выбрасываются дальше.
"""
try:
async with httpx.AsyncClient() as client:
response = await client.get(f"{super.params.url}/v1/openai/models", super.create_headers())
if response.status_code == 200:
json_data = response.json()
return [item['id'] for item in json_data.get('data', [])]
except httpx.RequestError as error:
print('Error fetching models:', error)
return []
def create_messages(self, prompt: str) -> List[dict]:
"""
Создает сообщения для LLM на основе переданного промпта и системного промпта (если он задан).
Args:
prompt (str): Пользовательский промпт.
Returns:
list[dict]: Список сообщений с ролями и содержимым.
"""
actual_prompt = self.apply_llm_template_to_prompt(prompt)
messages = []
if super.params.predict_params and super.params.predict_params.system_prompt:
messages.append({"role": "system", "content": super.params.predict_params.system_prompt})
messages.append({"role": "user", "content": actual_prompt})
return messages
def apply_llm_template_to_prompt(self, prompt: str) -> str:
"""
Применяет шаблон LLM к переданному промпту, если он задан.
Args:
prompt (str): Пользовательский промпт.
Returns:
str: Промпт с примененным шаблоном (или оригинальный, если шаблон отсутствует).
"""
actual_prompt = prompt
if super.params.template is not None:
actual_prompt = super.params.template.replace("{{PROMPT}}", actual_prompt)
return actual_prompt
async def tokenize(self, prompt: str) -> Optional[dict]:
raise NotImplementedError("This function is not supported.")
async def detokenize(self, tokens: List[int]) -> Optional[str]:
raise NotImplementedError("This function is not supported.")
async def create_request(self, prompt: str) -> dict:
"""
Создает запрос для предсказания на основе параметров LLM.
Args:
prompt (str): Промпт для запроса.
Returns:
dict: Словарь с параметрами для выполнения запроса.
"""
request = {
"stream": False,
"model": super.params.model,
}
predict_params = super.params.predict_params
if predict_params:
if predict_params.stop:
non_empty_stop = list(filter(lambda o: o != "", predict_params.stop))
if non_empty_stop:
request["stop"] = non_empty_stop
if predict_params.n_predict is not None:
request["max_tokens"] = int(predict_params.n_predict or 0)
request["temperature"] = float(predict_params.temperature or 0)
if predict_params.top_k is not None:
request["top_k"] = int(predict_params.top_k)
if predict_params.top_p is not None:
request["top_p"] = float(predict_params.top_p)
if predict_params.min_p is not None:
request["min_p"] = float(predict_params.min_p)
if predict_params.seed is not None:
request["seed"] = int(predict_params.seed)
if predict_params.n_keep is not None:
request["n_keep"] = int(predict_params.n_keep)
if predict_params.cache_prompt is not None:
request["cache_prompt"] = bool(predict_params.cache_prompt)
if predict_params.repeat_penalty is not None:
request["repetition_penalty"] = float(predict_params.repeat_penalty)
if predict_params.repeat_last_n is not None:
request["repeat_last_n"] = int(predict_params.repeat_last_n)
if predict_params.presence_penalty is not None:
request["presence_penalty"] = float(predict_params.presence_penalty)
if predict_params.frequency_penalty is not None:
request["frequency_penalty"] = float(predict_params.frequency_penalty)
request["messages"] = self.create_messages(prompt)
return request
async def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict:
raise NotImplementedError("This function is not supported.")
async def predict(self, prompt: str) -> str:
"""
Выполняет запрос к API и возвращает результат.
Args:
prompt (str): Входной текст для предсказания.
Returns:
str: Сгенерированный текст.
"""
async with httpx.AsyncClient() as client:
request = await self.create_request(prompt)
async with httpx.AsyncClient() as client:
response = client.post(f"{super.params.url}/v1/openai/chat/completions", super.create_headers(), json=request)
if response.status_code == 200:
return response.json()["choices"][0]["message"]["content"]
|