import gradio as gr import os from llama_index.core import Settings from documents_prep import load_json_documents, load_table_data, load_image_data, load_csv_chunks from utils import get_llm_model, get_embedding_model, get_reranker_model, answer_question from my_logging import log_message from index_retriever import create_vector_index, create_query_engine import sys from config import ( HF_REPO_ID, HF_TOKEN, DOWNLOAD_DIR, CHUNKS_FILENAME, JSON_FILES_DIR, TABLE_DATA_DIR, IMAGE_DATA_DIR, DEFAULT_MODEL, AVAILABLE_MODELS ) def initialize_system(repo_id, hf_token, download_dir, chunks_filename=None, json_files_dir=None, table_data_dir=None, image_data_dir=None, use_json_instead_csv=False): try: log_message("Инициализация системы") os.makedirs(download_dir, exist_ok=True) embed_model = get_embedding_model() llm = get_llm_model(DEFAULT_MODEL) reranker = get_reranker_model() Settings.embed_model = embed_model Settings.llm = llm all_documents = [] chunks_df = None if use_json_instead_csv and json_files_dir: log_message("Используем JSON файлы вместо CSV") json_documents = load_json_documents(repo_id, hf_token, json_files_dir, download_dir) all_documents.extend(json_documents) else: if chunks_filename: log_message("Загружаем данные из CSV") csv_documents, chunks_df = load_csv_chunks(repo_id, hf_token, chunks_filename, download_dir) all_documents.extend(csv_documents) if table_data_dir: log_message("Добавляю табличные данные") table_documents = load_table_data(repo_id, hf_token, table_data_dir) all_documents.extend(table_documents) if image_data_dir: log_message("Добавляю данные изображений") image_documents = load_image_data(repo_id, hf_token, image_data_dir) all_documents.extend(image_documents) log_message(f"Всего документов: {len(all_documents)}") vector_index = create_vector_index(all_documents) query_engine = create_query_engine(vector_index) log_message(f"Система успешно инициализирована") return query_engine, chunks_df, reranker, vector_index except Exception as e: log_message(f"Ошибка инициализации: {str(e)}") return None, None, None, None def switch_model(model_name, vector_index): from llama_index.core import Settings from index_retriever import create_query_engine try: log_message(f"Переключение на модель: {model_name}") new_llm = get_llm_model(model_name) Settings.llm = new_llm if vector_index is not None: new_query_engine = create_query_engine(vector_index) log_message(f"Модель успешно переключена на: {model_name}") return new_query_engine, f"✅ Модель переключена на: {model_name}" else: return None, "❌ Ошибка: система не инициализирована" except Exception as e: error_msg = f"Ошибка переключения модели: {str(e)}" log_message(error_msg) return None, f"❌ {error_msg}" def create_demo_interface(answer_question_func, switch_model_func, current_model): with gr.Blocks(title="AIEXP - AI Expert для нормативной документации", theme=gr.themes.Soft()) as demo: gr.Markdown(""" # AIEXP - Artificial Intelligence Expert ## Инструмент для работы с нормативной документацией """) with gr.Tab("🏠 Поиск по нормативным документам"): gr.Markdown("### Задайте вопрос по нормативной документации") with gr.Row(): with gr.Column(scale=2): model_dropdown = gr.Dropdown( choices=list(AVAILABLE_MODELS.keys()), value=current_model, label="🤖 Выберите языковую модель", info="Выберите модель для генерации ответов" ) with gr.Column(scale=1): switch_btn = gr.Button("🔄 Переключить модель", variant="secondary") model_status = gr.Textbox( value=f"Текущая модель: {current_model}", label="Статус модели", interactive=False ) with gr.Row(): with gr.Column(scale=3): question_input = gr.Textbox( label="Ваш вопрос к базе знаний", placeholder="Введите вопрос по нормативным документам...", lines=3 ) ask_btn = gr.Button("🔍 Найти ответ", variant="primary", size="lg") gr.Examples( examples=[ "О чем этот рисунок: ГОСТ Р 50.04.07-2022 Приложение Л. Л.1.5 Рисунок Л.2", "Л.9 Формула в ГОСТ Р 50.04.07 - 2022 что и о чем там?", "Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?", "Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?", "В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями?", ], inputs=question_input ) with gr.Row(): with gr.Column(scale=2): answer_output = gr.HTML( label="", value=f"
Здесь появится ответ на ваш вопрос...
Текущая модель: {current_model}
", ) with gr.Column(scale=1): sources_output = gr.HTML( label="", value="
Здесь появятся источники...
", ) switch_btn.click( fn=switch_model_func, inputs=[model_dropdown], outputs=[model_status] ) ask_btn.click( fn=answer_question_func, inputs=[question_input], outputs=[answer_output, sources_output] ) question_input.submit( fn=answer_question_func, inputs=[question_input], outputs=[answer_output, sources_output] ) return demo query_engine = None chunks_df = None reranker = None vector_index = None current_model = DEFAULT_MODEL def main_answer_question(question): global query_engine, reranker, current_model, chunks_df return answer_question(question, query_engine, reranker, current_model, chunks_df) def main_switch_model(model_name): global query_engine, vector_index, current_model new_query_engine, status_message = switch_model(model_name, vector_index) if new_query_engine: query_engine = new_query_engine current_model = model_name return status_message def main(): global query_engine, chunks_df, reranker, vector_index, current_model log_message("Запуск AIEXP - AI Expert для нормативной документации") query_engine, chunks_df, reranker, vector_index = initialize_system( repo_id=HF_REPO_ID, hf_token=HF_TOKEN, download_dir=DOWNLOAD_DIR, json_files_dir=JSON_FILES_DIR, table_data_dir=TABLE_DATA_DIR, image_data_dir=IMAGE_DATA_DIR, use_json_instead_csv=True, ) if query_engine: log_message("Запуск веб-интерфейса") demo = create_demo_interface( answer_question_func=main_answer_question, switch_model_func=main_switch_model, current_model=current_model ) demo.launch( server_name="0.0.0.0", server_port=7860, share=True, debug=False ) else: log_message("Невозможно запустить приложение из-за ошибки инициализации") sys.exit(1) if __name__ == "__main__": main()