diff --git a/Dockerfile b/Dockerfile index 2f885c2..8f93be0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -28,7 +28,7 @@ RUN pip install --upgrade pip && \ WORKDIR /app # Копируем только requirements.txt для кэширования слоя -COPY offline_packages/test.txt /app/requirements.txt +COPY offline_packages/requirements.txt /app/requirements.txt # Установка зависимостей (включая GPU-версии) RUN pip install --no-cache-dir -r /app/requirements.txt @@ -37,4 +37,4 @@ RUN pip install --no-cache-dir -r /app/requirements.txt COPY app /app # Запуск приложения -#CMD ["python", "main.py"] \ No newline at end of file +CMD ["python", "main.py"] \ No newline at end of file diff --git a/app/main.py b/app/main.py index 3128dc4..55b2ae7 100644 --- a/app/main.py +++ b/app/main.py @@ -1,77 +1,175 @@ +from typing import Any, Dict, List, Optional import gradio as gr -from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer import torch import threading +from loguru import logger +from langchain_chroma import Chroma +from langchain_huggingface import HuggingFaceEmbeddings +from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer -# Пути к модели -model_name = "/models/Qwen3-8B" +# Предполагается, что у тебя есть config.py с settings +from config import settings -# Загрузка токенизатора и модели (один раз при старте) -tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) -model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=torch.float16, - device_map="cuda", - trust_remote_code=True -) -# Отключим кэширование в истории, чтобы каждый запрос был независимым -def generate_response(message, history): - # Форматируем диалог: используем только текущую историю - prompt = "" - for human, assistant in history: - prompt += f"<|im_start|>user\n{human}<|im_end|>\n<|im_start|>assistant\n{assistant}<|im_end|>\n" - prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n" +class ChatWithAI: + def __init__(self, provider: str = "qwen3"): + self.provider = provider + self.embeddings = HuggingFaceEmbeddings( + model_name=settings.LM_MODEL_NAME, + model_kwargs={"device": "cuda"}, + encode_kwargs={"normalize_embeddings": True}, + ) - # Токенизация - inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + if provider == "qwen3": + model_name = getattr(settings, "LOCAL_LLM_NAME", "/models/Qwen3-4B") + logger.info(f"Загрузка локальной модели: {model_name}") - # Создаём уникальный streamer для каждого запроса - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt=True, - skip_special_tokens=True + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + self.model = AutoModelForCausalLM.from_pretrained( + model_name, + torch_dtype=torch.float32, + device_map="cuda", + ) + + # Streamer для потоковой генерации + self.streamer = TextIteratorStreamer( + self.tokenizer, + skip_prompt=True, + skip_special_tokens=True + ) + else: + raise ValueError(f"Неподдерживаемый провайдер: {provider}") + + self.chroma_db = Chroma( + persist_directory=settings.DOCS_CHROMA_PATH, + embedding_function=self.embeddings, + collection_name=settings.DOCS_COLLECTION_NAME, + ) + + def get_relevant_context(self, query: str, k: int = 3) -> List[Dict[str, Any]]: + """Получение релевантного контекста из базы данных.""" + try: + results = self.chroma_db.similarity_search(query, k=k) + return [ + { + "text": doc.page_content, + "metadata": doc.metadata, + } + for doc in results + ] + except Exception as e: + logger.error(f"Ошибка при получении контекста: {e}") + return [] + + def format_context(self, context: List[Dict[str, Any]]) -> str: + """Форматирование контекста для промпта.""" + formatted_context = [] + for item in context: + metadata_str = "\n".join(f"{k}: {v}" for k, v in item["metadata"].items()) + formatted_context.append( + f"Текст: {item['text']}\nМетаданные:\n{metadata_str}\n" + ) + return "\n---\n".join(formatted_context) + + def generate_response_stream(self, query: str): + """Генерация ответа с потоковой передачей токенов.""" + try: + logger.info(f"Пользовательский запрос: {query}") + context = self.get_relevant_context(query) + if not context: + yield "Извините, не удалось найти релевантный контекст для ответа." + return + + formatted_context = self.format_context(context) + + messages = [ + { + "role": "system", + "content": """Ты — внутренний менеджер компании Mazban. Отвечаешь по делу без лишних вступлений. + +Правила: +1. Сразу переходи к сути, без фраз типа "На основе контекста" +2. Используй только факты. Если точных данных нет — отвечай общими фразами об Marzban, но не придумывай конкретику +3. Используй обычный текст без форматирования +4. Включай ссылки только если они есть в контексте +5. Говори от первого лица множественного числа: "Мы предоставляем", "У нас есть" +6. При упоминании файлов делай это естественно, например: "Я прикреплю инструкцию, где подробно описаны шаги" +7. На приветствия отвечай доброжелательно, на негатив — с легким юмором +8. Можешь при ответах использовать общую информацию из открытых источников по Marzban, но опирайся на контекст +9. Если пользователь спрашивает о ценах, планах или технических характеристиках — давай конкретные ответы из контекста +10. При технических вопросах предлагай практические решения + +Персонализируй ответы, упоминая имя клиента если оно есть в контексте. Будь краток, информативен и полезен.""", + }, + { + "role": "user", + "content": f"Вопрос: {query}\nКонтекст: {formatted_context}", + }, + ] + + # Применяем шаблон чата + prompt = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False + ) + + # Подготавливаем вход + inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda") + + # Очищаем streamer и запускаем генерацию в отдельном потоке + self.streamer = TextIteratorStreamer( + self.tokenizer, + skip_prompt=True, + skip_special_tokens=True + ) + + generate_kwargs = { + "input_ids": inputs["input_ids"], + "max_new_tokens": 512, + "temperature": 0.7, + "do_sample": True, + "top_p": 0.9, + "pad_token_id": self.tokenizer.eos_token_id, + "streamer": self.streamer, + } + + thread = threading.Thread(target=self.model.generate, kwargs=generate_kwargs) + thread.start() + + # Потоковая передача токенов + buffer = "" + for token in self.streamer: + buffer += token + yield buffer # Отправляем частичный ответ + + except Exception as e: + logger.error(f"Ошибка при генерации ответа: {e}") + yield "Произошла ошибка при генерации ответа." + + +# === Gradio интерфейс === +def main(): + chat = ChatWithAI(provider="qwen3") + + def respond(message, history): + # Генерируем ответ по частям + for token in chat.generate_response_stream(message): + yield token + + demo = gr.ChatInterface( + fn=respond, + title="Помощник настройки Marzban", + description="Задайте вопрос — получите ответ от внутреннего менеджера.", + examples=[ + "Как подключить Marzban?", + "Как настроить telegram бота?", + "Что такое marzban?" + ], ) + demo.launch(server_name="0.0.0.0", server_port=8080, share=False) - # Параметры генерации - generation_kwargs = { - "input_ids": inputs["input_ids"], - "max_new_tokens": 1024, - "temperature": 0.6, - "top_p": 0.9, - "do_sample": True, - "pad_token_id": tokenizer.eos_token_id, - "streamer": streamer, - } - - # Запускаем генерацию в отдельном потоке - thread = threading.Thread(target=model.generate, kwargs=generation_kwargs) - thread.start() - - # Постепенно возвращаем результат - buffer = "" - for new_text in streamer: - buffer += new_text - yield buffer.strip() - -# Создаем интерфейс -demo = gr.ChatInterface( - fn=generate_response, - title="Qwen3-4B-Base Chat", - description="Общайтесь с моделью Qwen3-4B-Base в режиме реального времени с потоковой генерацией", - examples=[ - "Объясни, как работает квантование AWQ?", - "Напиши стихотворение про ИИ", - "Какие преимущества у Qwen3 перед предыдущими версиями?" - ], - theme="soft", -) if __name__ == "__main__": - # ВАЖНО: используем .queue() для поддержки асинхронной обработки - demo.queue(max_size=20, default_concurrency_limit=10).launch( - server_port=8080, - server_name="0.0.0.0", - share=False, - # Можно добавить: max_batch_size=1, concurrency_count=4 - ) \ No newline at end of file + main() \ No newline at end of file diff --git a/app/old_without_db.py b/app/old_without_db.py new file mode 100644 index 0000000..3128dc4 --- /dev/null +++ b/app/old_without_db.py @@ -0,0 +1,77 @@ +import gradio as gr +from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer +import torch +import threading + +# Пути к модели +model_name = "/models/Qwen3-8B" + +# Загрузка токенизатора и модели (один раз при старте) +tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) +model = AutoModelForCausalLM.from_pretrained( + model_name, + torch_dtype=torch.float16, + device_map="cuda", + trust_remote_code=True +) + +# Отключим кэширование в истории, чтобы каждый запрос был независимым +def generate_response(message, history): + # Форматируем диалог: используем только текущую историю + prompt = "" + for human, assistant in history: + prompt += f"<|im_start|>user\n{human}<|im_end|>\n<|im_start|>assistant\n{assistant}<|im_end|>\n" + prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n" + + # Токенизация + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + + # Создаём уникальный streamer для каждого запроса + streamer = TextIteratorStreamer( + tokenizer, + skip_prompt=True, + skip_special_tokens=True + ) + + # Параметры генерации + generation_kwargs = { + "input_ids": inputs["input_ids"], + "max_new_tokens": 1024, + "temperature": 0.6, + "top_p": 0.9, + "do_sample": True, + "pad_token_id": tokenizer.eos_token_id, + "streamer": streamer, + } + + # Запускаем генерацию в отдельном потоке + thread = threading.Thread(target=model.generate, kwargs=generation_kwargs) + thread.start() + + # Постепенно возвращаем результат + buffer = "" + for new_text in streamer: + buffer += new_text + yield buffer.strip() + +# Создаем интерфейс +demo = gr.ChatInterface( + fn=generate_response, + title="Qwen3-4B-Base Chat", + description="Общайтесь с моделью Qwen3-4B-Base в режиме реального времени с потоковой генерацией", + examples=[ + "Объясни, как работает квантование AWQ?", + "Напиши стихотворение про ИИ", + "Какие преимущества у Qwen3 перед предыдущими версиями?" + ], + theme="soft", +) + +if __name__ == "__main__": + # ВАЖНО: используем .queue() для поддержки асинхронной обработки + demo.queue(max_size=20, default_concurrency_limit=10).launch( + server_port=8080, + server_name="0.0.0.0", + share=False, + # Можно добавить: max_batch_size=1, concurrency_count=4 + ) \ No newline at end of file diff --git a/app/test.py b/app/test.py deleted file mode 100644 index aeaffb2..0000000 --- a/app/test.py +++ /dev/null @@ -1,161 +0,0 @@ -from typing import Any, Dict, List, Literal, Optional - -import torch -from config import settings -from langchain_chroma import Chroma -from langchain_huggingface import HuggingFaceEmbeddings -from loguru import logger -from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline - - -class ChatWithAI: - def __init__(self, provider: "qwen3"): - self.provider = provider - self.embeddings = HuggingFaceEmbeddings( - model_name=settings.LM_MODEL_NAME, - model_kwargs={"device": "cuda"}, - encode_kwargs={"normalize_embeddings": True}, - ) - - if provider == "qwen3": - model_name = getattr(settings, "LOCAL_LLM_NAME", "/models/Qwen3-4B") - - logger.info(f"Загрузка локальной модели: {model_name}") - tokenizer = AutoTokenizer.from_pretrained(model_name) - model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=torch.float32, - device_map="cuda", - ) - - #Создаём text-generation pipeline - self.llm = pipeline( - "text-generation", - model=model, - tokenizer=tokenizer, - device_map="cuda", # 0 = GPU, -1 = CPU - temperature=0.7, - max_new_tokens=512, - do_sample=True, - pad_token_id=tokenizer.eos_token_id, - ) - else: - raise ValueError(f"Неподдерживаемый провайдер: {provider}") - - self.chroma_db = Chroma( - persist_directory=settings.DOCS_CHROMA_PATH, - embedding_function=self.embeddings, - collection_name=settings.DOCS_COLLECTION_NAME, - ) - - def get_relevant_context(self, query: str, k: int = 3) -> List[Dict[str, Any]]: - """Получение релевантного контекста из базы данных.""" - try: - results = self.chroma_db.similarity_search(query, k=k) - return [ - { - "text": doc.page_content, - "metadata": doc.metadata, - } - for doc in results - ] - except Exception as e: - logger.error(f"Ошибка при получении контекста: {e}") - return [] - - def format_context(self, context: List[Dict[str, Any]]) -> str: - """Форматирование контекста для промпта.""" - formatted_context = [] - for item in context: - metadata_str = "\n".join(f"{k}: {v}" for k, v in item["metadata"].items()) - formatted_context.append( - f"Текст: {item['text']}\nМетаданные:\n{metadata_str}\n" - ) - return "\n---\n".join(formatted_context) - - def generate_response(self, query: str) -> Optional[str]: - """Генерация ответа на основе запроса и контекста.""" - try: - context = self.get_relevant_context(query) - if not context: - return "Извините, не удалось найти релевантный контекст для ответа." - - formatted_context = self.format_context(context) - - messages = [ - { - "role": "system", - "content": """Ты — внутренний менеджер компании Amvera Cloud. Отвечаешь по делу без лишних вступлений. - -Правила: -1. Сразу переходи к сути, без фраз типа "На основе контекста" -2. Используй только факты. Если точных данных нет — отвечай общими фразами об Marzban, но не придумывай конкретику -3. Используй обычный текст без форматирования -4. Включай ссылки только если они есть в контексте -5. Говори от первого лица множественного числа: "Мы предоставляем", "У нас есть" -6. При упоминании файлов делай это естественно, например: "Я прикреплю инструкцию, где подробно описаны шаги" -7. На приветствия отвечай доброжелательно, на негатив — с легким юмором -8. Можешь при ответах использовать общую информацию из открытых источников по Marzban, но опирайся на контекст -9. Если пользователь спрашивает о ценах, планах или технических характеристиках — давай конкретные ответы из контекста -10. При технических вопросах предлагай практические решения - -Персонализируй ответы, упоминая имя клиента если оно есть в контексте. Будь краток, информативен и полезен.""", - }, - { - "role": "user", - "content": f"Вопрос: {query}\nКонтекст: {formatted_context}", - }, - ] - # Генерация через transformers - if self.provider == "qwen3": - # Используем токенизатор модели для форматирования чата - tokenizer = self.llm.tokenizer - model = self.llm.model - - # Применяем chat template (поддерживается в современных моделях: Zephyr, Llama3, Qwen и т.д.) - prompt = tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True, enable_thinking=False - ) - - # Генерация - outputs = self.llm( - prompt, - max_new_tokens=512, - temperature=0.7, - do_sample=True, - top_p=0.9, - pad_token_id=tokenizer.eos_token_id, - ) - response_text = outputs[0]["generated_text"] - - # Убираем входной промпт, оставляем только ответ - if prompt in response_text: - response_text = response_text[len(prompt):].strip() - - return response_text - - else: - # Остальные провайдеры (deepseek, openai) используют langchain - response = self.llm.invoke(messages) - if hasattr(response, "content"): - return str(response.content) - return str(response).strip() - except Exception as e: - logger.error(f"Ошибка при генерации ответа: {e}") - return "Произошла ошибка при генерации ответа." - - -if __name__ == "__main__": - chat = ChatWithAI(provider="qwen3") - print("\n=== Чат с ИИ ===\n") - - while True: - query = input("Вы: ") - if query.lower() == "выход": - print("\nДо свидания!") - break - - print("\nИИ печатает...", end="\r") - response = chat.generate_response(query) - print(" " * 20, end="\r") # Очищаем "ИИ печатает..." - print(f"ИИ: {response}\n") \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 05b72b4..47966b0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -20,7 +20,7 @@ services: - ./cache/cache:/root/.cache - ./cache/site-packages:/usr/local/lib/python3.10/site-packages - ./offline_packages:/offline_packages - entrypoint: sleep 1000000 #./app/entrypoint.sh + # entrypoint: sleep 1000000 #./app/entrypoint.sh ports: - "8080:8080" networks: