43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
# app/worker.py
|
|
import time
|
|
from app.model_loader import ModelLoader
|
|
from app.inference_service import InferenceService
|
|
from app.redis_client import RedisClient
|
|
from config import BASE_MODEL, ADAPTER_DIR, HF_TOKEN, REDIS_HOST, REDIS_PORT, TEXT_RESULT_CHANNEL, TEXT_TASK_CHANNEL, BATCH_SIZE, WAIT_TIMEOUT
|
|
|
|
def main():
|
|
model_loader = ModelLoader(BASE_MODEL, ADAPTER_DIR, HF_TOKEN)
|
|
model_loader.load_model()
|
|
model_loader.load_tokenizer()
|
|
inference_service = InferenceService(model_loader)
|
|
|
|
redis_client = RedisClient(
|
|
host=REDIS_HOST,
|
|
port=REDIS_PORT,
|
|
task_channel=TEXT_TASK_CHANNEL,
|
|
result_channel=TEXT_RESULT_CHANNEL
|
|
)
|
|
|
|
print("Worker запущен, ожидаем задачи...")
|
|
|
|
while True:
|
|
tasks = redis_client.get_tasks(BATCH_SIZE, wait_timeout=WAIT_TIMEOUT)
|
|
if not tasks:
|
|
time.sleep(0.5)
|
|
continue
|
|
|
|
texts = [task.text for task in tasks]
|
|
responses = inference_service.generate_batch(texts)
|
|
for task, response in zip(tasks, responses):
|
|
result = {
|
|
"chat_id": task.chat_id,
|
|
"user_id": task.user_id,
|
|
"message_id": task.message_id,
|
|
"text": response
|
|
}
|
|
redis_client.publish_result(result)
|
|
print(f"Обработана задача {task.message_id}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|