add vector db
This commit is contained in:
6
.gitignore
vendored
6
.gitignore
vendored
@@ -1,4 +1,8 @@
|
||||
/app/.gradio
|
||||
/app/core
|
||||
/app/shop_chroma_db
|
||||
/app/amvera_data/chroma_db
|
||||
/app/amvera_data/parsed_json/
|
||||
/app/__pycache__
|
||||
/models
|
||||
/cache
|
||||
/cache
|
||||
|
||||
@@ -28,7 +28,7 @@ RUN pip install --upgrade pip && \
|
||||
WORKDIR /app
|
||||
|
||||
# Копируем только requirements.txt для кэширования слоя
|
||||
COPY offline_packages/requirements.txt /app/requirements.txt
|
||||
COPY offline_packages/test.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"]
|
||||
#CMD ["python", "main.py"]
|
||||
120
app/amvera_data/generate_json_files.py
Normal file
120
app/amvera_data/generate_json_files.py
Normal file
@@ -0,0 +1,120 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import string
|
||||
import sys
|
||||
from typing import Any, Dict, List
|
||||
|
||||
# Добавляем путь к родительской директории для импорта config
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from config import settings
|
||||
from loguru import logger
|
||||
|
||||
# Константы
|
||||
HEADER_PATTERN = re.compile(r"^(#+)\s(.+)")
|
||||
PUNCTUATION_PATTERN = re.compile(f"[{re.escape(string.punctuation)}]")
|
||||
WHITESPACE_PATTERN = re.compile(r"\s+")
|
||||
|
||||
|
||||
def normalize_text(text: str) -> str:
|
||||
"""Нормализация текста: удаление знаков препинания и специальных символов."""
|
||||
if not isinstance(text, str):
|
||||
raise ValueError("Входной текст должен быть строкой")
|
||||
|
||||
# Удаление знаков препинания
|
||||
text = PUNCTUATION_PATTERN.sub(" ", text)
|
||||
# Удаление переносов строк и лишних пробелов
|
||||
text = WHITESPACE_PATTERN.sub(" ", text)
|
||||
# Приведение к нижнему регистру
|
||||
return text.lower().strip()
|
||||
|
||||
|
||||
def parse_markdown(md_path: str) -> Dict[str, Any]:
|
||||
"""Парсинг markdown файла и создание структурированных данных."""
|
||||
if not os.path.exists(md_path):
|
||||
raise FileNotFoundError(f"Файл {md_path} не найден")
|
||||
|
||||
try:
|
||||
with open(md_path, "r", encoding="utf-8") as file:
|
||||
content = file.read()
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при чтении файла {md_path}: {e}")
|
||||
raise
|
||||
|
||||
sections: List[str] = []
|
||||
section_titles: List[str] = []
|
||||
current_section: str | None = None
|
||||
current_content: List[str] = []
|
||||
|
||||
for line in content.splitlines():
|
||||
section_match = HEADER_PATTERN.match(line)
|
||||
|
||||
if section_match:
|
||||
if current_section:
|
||||
sections.append("\n".join(current_content).strip())
|
||||
section_titles.append(current_section)
|
||||
current_content = []
|
||||
current_section = section_match.group(2)
|
||||
current_content.append(current_section)
|
||||
else:
|
||||
current_content.append(line)
|
||||
|
||||
if current_section:
|
||||
sections.append("\n".join(current_content).strip())
|
||||
section_titles.append(current_section)
|
||||
|
||||
# Нормализация текста для векторной базы данных
|
||||
normalized_sections = [normalize_text(section) for section in sections]
|
||||
full_text = " ".join(normalized_sections)
|
||||
|
||||
# Создаем структуру метаданных
|
||||
metadata = {
|
||||
"file_name": os.path.basename(md_path),
|
||||
"section_count": len(section_titles),
|
||||
}
|
||||
|
||||
# Добавляем заголовки как отдельные поля
|
||||
for i, title in enumerate(section_titles):
|
||||
metadata[f"section_{i+1}"] = title
|
||||
|
||||
return {"text": full_text, "metadata": metadata}
|
||||
|
||||
|
||||
def process_all_markdown(input_folder: str, output_folder: str) -> None:
|
||||
"""Обработка всех markdown файлов в директории."""
|
||||
if not os.path.exists(input_folder):
|
||||
raise FileNotFoundError(f"Входная директория {input_folder} не найдена")
|
||||
|
||||
try:
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при создании выходной директории: {e}")
|
||||
raise
|
||||
|
||||
for root, _, files in os.walk(input_folder):
|
||||
for file_name in files:
|
||||
if file_name.endswith(".md"):
|
||||
try:
|
||||
md_path = os.path.join(root, file_name)
|
||||
output_path = os.path.join(
|
||||
output_folder, file_name.replace(".md", ".json")
|
||||
)
|
||||
parsed_data = parse_markdown(md_path)
|
||||
|
||||
with open(output_path, "w", encoding="utf-8") as file:
|
||||
json.dump(parsed_data, file, ensure_ascii=False, indent=4)
|
||||
logger.info(f"Результат сохранен в {output_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при обработке файла {file_name}: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
process_all_markdown(
|
||||
input_folder=settings.DOCS_AMVERA_PATH,
|
||||
output_folder=settings.PARSED_JSON_PATH,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Критическая ошибка: {e}")
|
||||
sys.exit(1)
|
||||
115
app/amvera_data/generated_amvera_db.py
Normal file
115
app/amvera_data/generated_amvera_db.py
Normal file
@@ -0,0 +1,115 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import torch
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from loguru import logger
|
||||
|
||||
# Добавляем путь к родительской директории для импорта config
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from config import settings
|
||||
|
||||
|
||||
def load_json_files(directory: str) -> List[Dict[str, Any]]:
|
||||
"""Загрузка всех JSON файлов из указанной директории."""
|
||||
documents = []
|
||||
|
||||
try:
|
||||
if not os.path.exists(directory):
|
||||
logger.error(f"Директория {directory} не существует")
|
||||
return documents
|
||||
|
||||
for filename in os.listdir(directory):
|
||||
if filename.endswith(".json"):
|
||||
file_path = os.path.join(directory, filename)
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as file:
|
||||
data = json.load(file)
|
||||
documents.append(
|
||||
{"text": data["text"], "metadata": data["metadata"]}
|
||||
)
|
||||
logger.info(f"Загружен файл: {filename}")
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при чтении файла {filename}: {e}")
|
||||
|
||||
logger.success(f"Загружено {len(documents)} JSON файлов")
|
||||
return documents
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при загрузке JSON файлов: {e}")
|
||||
return documents
|
||||
|
||||
|
||||
def split_text_into_chunks(text: str, metadata: Dict[str, Any]) -> List[Any]:
|
||||
"""Разделение текста на чанки с сохранением метаданных."""
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=settings.MAX_CHUNK_SIZE,
|
||||
chunk_overlap=settings.CHUNK_OVERLAP,
|
||||
length_function=len,
|
||||
is_separator_regex=False,
|
||||
)
|
||||
|
||||
chunks = text_splitter.create_documents(texts=[text], metadatas=[metadata])
|
||||
return chunks
|
||||
|
||||
|
||||
def generate_chroma_db() -> Optional[Chroma]:
|
||||
"""Инициализация ChromaDB с данными из JSON файлов."""
|
||||
try:
|
||||
# Создаем директорию для хранения базы данных, если она не существует
|
||||
os.makedirs(settings.AMVERA_CHROMA_PATH, exist_ok=True)
|
||||
|
||||
# Загружаем JSON файлы
|
||||
documents = load_json_files(settings.PARSED_JSON_PATH)
|
||||
|
||||
if not documents:
|
||||
logger.warning("Нет документов для добавления в базу данных")
|
||||
return None
|
||||
|
||||
# Инициализируем модель эмбеддингов
|
||||
embeddings = HuggingFaceEmbeddings(
|
||||
model_name=settings.LM_MODEL_NAME,
|
||||
model_kwargs={"device": "cuda" if torch.cuda.is_available() else "cpu"},
|
||||
encode_kwargs={"normalize_embeddings": True},
|
||||
)
|
||||
|
||||
# Подготавливаем данные для Chroma
|
||||
all_chunks = []
|
||||
for i, doc in enumerate(documents):
|
||||
chunks = split_text_into_chunks(doc["text"], doc["metadata"])
|
||||
all_chunks.extend(chunks)
|
||||
logger.info(
|
||||
f"Документ {i+1}/{len(documents)} разбит на {len(chunks)} чанков"
|
||||
)
|
||||
|
||||
# Создаем векторное хранилище
|
||||
texts = [chunk.page_content for chunk in all_chunks]
|
||||
metadatas = [chunk.metadata for chunk in all_chunks]
|
||||
ids = [f"doc_{i}" for i in range(len(all_chunks))]
|
||||
|
||||
chroma_db = Chroma.from_texts(
|
||||
texts=texts,
|
||||
embedding=embeddings,
|
||||
ids=ids,
|
||||
metadatas=metadatas,
|
||||
persist_directory=settings.AMVERA_CHROMA_PATH,
|
||||
collection_name=settings.AMVERA_COLLECTION_NAME,
|
||||
collection_metadata={
|
||||
"hnsw:space": "cosine",
|
||||
},
|
||||
)
|
||||
|
||||
logger.success(
|
||||
f"База Chroma инициализирована, добавлено {len(all_chunks)} чанков из {len(documents)} документов"
|
||||
)
|
||||
return chroma_db
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка инициализации Chroma: {e}")
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
generate_chroma_db()
|
||||
20
app/config.py
Normal file
20
app/config.py
Normal file
@@ -0,0 +1,20 @@
|
||||
import os
|
||||
|
||||
from pydantic import SecretStr
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class Config(BaseSettings):
|
||||
BASE_DIR: str = os.path.abspath(os.path.join(os.path.dirname(__file__)))
|
||||
DOCS_AMVERA_PATH: str = os.path.join(BASE_DIR, "amvera_data", "docs_amvera")
|
||||
PARSED_JSON_PATH: str = os.path.join(BASE_DIR, "amvera_data", "parsed_json")
|
||||
AMVERA_CHROMA_PATH: str = os.path.join(BASE_DIR, "amvera_data", "chroma_db")
|
||||
AMVERA_COLLECTION_NAME: str = "amvera_docs"
|
||||
MAX_CHUNK_SIZE: int = 512
|
||||
CHUNK_OVERLAP: int = 50
|
||||
LM_MODEL_NAME: str = "/models/paraphrase-multilingual-MiniLM-L12-v2"
|
||||
LOCAL_LLM_NAME: str = "/models/Qwen3-4B-Base"
|
||||
QWEN_MODEL_NAME: str = "qwen3-chat"
|
||||
|
||||
|
||||
settings = Config() # type: ignore
|
||||
161
app/test.py
Normal file
161
app/test.py
Normal file
@@ -0,0 +1,161 @@
|
||||
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-Base")
|
||||
|
||||
logger.info(f"Загрузка локальной модели: {model_name}")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=torch.float16,
|
||||
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.AMVERA_CHROMA_PATH,
|
||||
embedding_function=self.embeddings,
|
||||
collection_name=settings.AMVERA_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. Используй только факты. Если точных данных нет — отвечай общими фразами об Amvera Cloud, но не придумывай конкретику
|
||||
3. Используй обычный текст без форматирования
|
||||
4. Включай ссылки только если они есть в контексте
|
||||
5. Говори от первого лица множественного числа: "Мы предоставляем", "У нас есть"
|
||||
6. При упоминании файлов делай это естественно, например: "Я прикреплю инструкцию, где подробно описаны шаги"
|
||||
7. На приветствия отвечай доброжелательно, на негатив — с легким юмором
|
||||
8. Можешь при ответах использовать общую информацию из открытых источников по Amvera Cloud, но опирайся на контекст
|
||||
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
|
||||
)
|
||||
|
||||
# Генерация
|
||||
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")
|
||||
@@ -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:
|
||||
|
||||
@@ -1,5 +1,13 @@
|
||||
torch>=2.6.0
|
||||
|
||||
transformers>=4.50.0
|
||||
gradio>=4.29.0
|
||||
safetensors>=0.4.3
|
||||
accelerate>=0.29.3
|
||||
accelerate>=1.9.0
|
||||
langchain-huggingface==0.1.2
|
||||
torch>=2.7.1
|
||||
loguru==0.7.3
|
||||
chromadb==0.6.3
|
||||
sentence-transformers==3.4.1
|
||||
langchain-chroma==0.2.2
|
||||
pydantic-settings==2.8.1
|
||||
langchain-text-splitters==0.3.7
|
||||
10
offline_packages/test.txt
Normal file
10
offline_packages/test.txt
Normal file
@@ -0,0 +1,10 @@
|
||||
langchain-huggingface==0.1.2
|
||||
torch>=2.7.1
|
||||
loguru==0.7.3
|
||||
chromadb==0.6.3
|
||||
sentence-transformers==3.4.1
|
||||
langchain-chroma==0.2.2
|
||||
pydantic-settings==2.8.1
|
||||
langchain-text-splitters==0.3.7
|
||||
langchain-qwen3==0.1.0
|
||||
langchain-openai==0.3.11
|
||||
Reference in New Issue
Block a user