app.py

import os
import torch
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.llms import HuggingFacePipeline
from transformers import pipeline

# Убедись, что все нужные модели и данные сохранены в этой папке
PERSIST_DIRECTORY = "db"

# Создание эмбеддингов
def create_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# Загрузка векторного хранилища
def load_vectorstore():
    embeddings = create_embeddings()
    return Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)

# Создание пайплайна LLM (используем GPU, если доступен)
def create_llm_pipeline():
    return pipeline(
        "text-generation",
        model="IlyaGusev/saiga2_7b_lora",
        device=0 if torch.cuda.is_available() else -1,
        max_new_tokens=512,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7
    )

# Обёртка LLM для использования с LangChain
def create_llm():
    pipe = create_llm_pipeline()
    return HuggingFacePipeline(pipeline=pipe)

# Сборка цепочки QA
def build_chain():
    vectorstore = load_vectorstore()
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    llm = create_llm()
    return RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        return_source_documents=True
    )

# Создаём QA цепочку
qa_chain = build_chain()

# Пример запроса (можешь удалить или адаптировать)
if __name__ == "__main__":
    question = "Кто такой Виктор Цой?"
    result = qa_chain(question)
    print("\nОтвет:\n", result["result"])