app.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import os
import re

# Конфигурация модели
MODEL_NAME = "Qwen/Qwen2.5-Coder-14B-Instruct-GPTQ-Int4"
DEVICE = "cpu"  # Запуск на CPU

# Глобальная загрузка модели (один раз при запуске)
def load_model():
    print("🔄 Загружаем токенизатор...")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True
    )
    
    # Добавляем pad token если его нет
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    print("🔄 Загружаем GPTQ-модель Qwen2.5-Coder-14B...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float32,  # Используем float32 для CPU
        device_map="cpu",           # Явно указываем CPU
        trust_remote_code=True
    )
    model.eval()  # Переводим модель в режим оценки
    return tokenizer, model

# Загружаем модель один раз при старте
try:
    tokenizer, model = load_model()
    print("✅ Qwen2.5-Coder-14B-Instruct-GPTQ-Int4 успешно загружена!")
except Exception as e:
    print(f"❌ Ошибка загрузки модели: {e}")
    tokenizer, model = None, None

def read_file_content(file_path):
    """Читает содержимое файла с обработкой различных кодировок"""
    try:
        # Пробуем разные кодировки
        encodings = ['utf-8', 'cp1251', 'iso-8859-1', 'windows-1251']
        
        for encoding in encodings:
            try:
                with open(file_path, 'r', encoding=encoding) as f:
                    content = f.read()
                return content
            except UnicodeDecodeError:
                continue
        
        # Если текстовые кодировки не работают, пробуем бинарный режим
        if file_path.endswith(('.py', '.txt', '.js', '.html', '.css', '.json', '.md')):
            with open(file_path, 'rb') as f:
                content = f.read()
            return content.decode('utf-8', errors='replace')
        else:
            return f"Файл {os.path.basename(file_path)} не является текстовым файлом"
            
    except Exception as e:
        return f"Ошибка чтения файла: {str(e)}"

def extract_search_terms(prompt):
    """Извлекает поисковые термины из промпта в формате [поиск: ...]"""
    search_pattern = r'\[поиск:\s*(.*?)\]'
    matches = re.findall(search_pattern, prompt, re.IGNORECASE)
    
    if matches:
        # Удаляем поисковую часть из промпта
        clean_prompt = re.sub(search_pattern, '', prompt).strip()
        return clean_prompt, matches[0].strip()
    
    return prompt, None

def generate_code_with_context(prompt, files, max_length=1024, temperature=0.7, top_p=0.9):
    """
    Генерирует код на основе промпта пользователя с учетом загруженных файлов.
    
    Эта функция автоматически станет доступна как MCP-инструмент для других приложений.
    
    Args:
        prompt (str): Запрос пользователя, может содержать [поиск: термин]
        files (list): Список загруженных файлов для анализа
        max_length (int): Максимальная длина ответа в токенах
        temperature (float): Параметр температуры для генерации
        top_p (float): Параметр top-p для генерации
        
    Returns:
        str: Сгенерированный код или сообщение об ошибке
    """
    if model is None or tokenizer is None:
        return "❌ Ошибка: модель не загружена. Проверьте:\n- Подключение к интернету\n- Достаточно ли оперативной памяти (рекомендуется 16+ ГБ)\n- Установлены ли зависимости: `pip install auto-gptq optimum`"
    
    try:
        # Извлекаем поисковые термины из промпта
        clean_prompt, search_term = extract_search_terms(prompt)
        
        # Обрабатываем загруженные файлы
        file_contexts = []
        
        if files:
            for file_info in files:
                if hasattr(file_info, 'name'):
                    file_path = file_info.name
                else:
                    file_path = file_info
                
                content = read_file_content(file_path)
                filename = os.path.basename(file_path)
                file_contexts.append(f"Файл: {filename}\n```\n{content[:2000]}\n```")
        
        # Формируем финальный промпт
        final_prompt = clean_prompt
        if file_contexts:
            files_context = "\n\n".join(file_contexts)
            final_prompt = f"""Контекст из загруженных файлов:
{files_context}

Запрос: {clean_prompt}"""
        
        # Форматируем сообщение для модели в соответствии с официальным форматом Qwen2.5 :cite[1]
        messages = [
            {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
            {"role": "user", "content": final_prompt}
        ]
        
        # Применяем шаблон чата :cite[1]
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        # Токенизируем входные данные с ограничением длины
        inputs = tokenizer(
            text, 
            return_tensors="pt", 
            truncation=True, 
            max_length=2048,
            padding=True,
            return_attention_mask=True
        )
    
        # Генерация с использованием обновленного кода и attention_mask
        with torch.no_grad():
            generated_ids = model.generate(
                inputs.input_ids,
                attention_mask=inputs.attention_mask,  # Добавляем attention_mask
                max_new_tokens=max_length,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1,
                no_repeat_ngram_size=3
            )
            
            # Декодируем результат
            generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
            
            # Убираем оригинальный промпт из ответа
            if generated_text.startswith(text):
                response = generated_text[len(text):].strip()
            else:
                response = generated_text.strip()
            
            return response
        
    except Exception as e:
        return f"❌ Ошибка при генерации кода: {str(e)}"

# Создаем интерфейс Gradio
with gr.Blocks(title="Qwen2.5-Coder-14B with MCP", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🚀 Qwen2.5-Coder-14B-Instruct-GPTQ-Int4 + MCP
    **Профессиональный генератор кода с поддержкой Model Context Protocol**
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                lines=4,
                placeholder="""Введите ваш запрос... Примеры:
- "Напиши функцию для быстрой сортировки на Python"
- "Создай REST API на FastAPI для управления пользователями"
- "Найди и исправь ошибку в загруженном коде [поиск: syntax error]" """,
                label="Запрос на генерацию кода",
                info="Используйте [поиск: ...] для поиска в файлах"
            )
            
            with gr.Accordion("📁 Загрузка файлов для анализа", open=True):
                file_input = gr.File(
                    label="Загрузите файлы для анализа",
                    file_count="multiple",
                    file_types=[".txt", ".py", ".js", ".html", ".css", ".json", ".md", ".java", ".cpp", ".c"],
                    type="filepath"
                )
            
            with gr.Accordion("⚙️ Параметры генерации", open=False):
                max_length_slider = gr.Slider(
                    minimum=100, maximum=2048, value=512,
                    label="Максимальная длина ответа (токены)"
                )
                temperature_slider = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.7,
                    label="Температура (креативность)"
                )
                top_p_slider = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.9,
                    label="Top-p (вероятностный отбор)"
                )
            
            generate_btn = gr.Button("🚀 Сгенерировать код", variant="primary")
            
        with gr.Column():
            response_output = gr.Textbox(
                label="Сгенерированный код",
                lines=18,
                show_copy_button=True
            )
    
    # Добавляем информацию о MCP
    with gr.Accordion("🔗 MCP Сервер - Подключение к другим приложениям", open=True):
        gr.Markdown("""
        **MCP (Model Context Protocol) сервер активирован!**
        
        Ваш генератор кода теперь доступен как MCP-инструмент для:
        - Claude Desktop
        - Cursor 
        - Cline
        - Других MCP-клиентов
        
        **URL для подключения:**
        - Основной MCP URL: `http://localhost:7860/gradio_api/mcp/`
        - SSE URL: `http://localhost:7860/gradio_api/mcp/sse`
        
        **Для подключения к Claude Desktop** добавьте в настройки (`claude_desktop_config.json`):
        ```json
        {
          "mcpServers": {
            "qwen-coder-generator": {
              "url": "http://localhost:7860/gradio_api/mcp/sse"
            }
          }
        }
        ```
        """)
    
    # Информация о модели
    with gr.Accordion("ℹ️ О модели Qwen2.5-Coder-14B", open=False):
        gr.Markdown("""
        **Qwen2.5-Coder-14B-Instruct** - это специализированная модель для программирования :cite[6]:
        - **Параметры**: 14.7 миллиардов (квантованные в INT4)
        - **Специализация**: Генерация кода, исправление ошибок, код-ризонинг
        - **Контекст**: До 128K токенов :cite[1]
        - **Языки программирования**: Поддержка 40+ языков :cite[6]
        - **Память**: ~4-6 ГБ RAM (благодаря GPTQ-квантованию)
        
        **Улучшения Qwen2.5 по сравнению с Qwen2** :cite[1]:
        - Значительно больше знаний и улучшенные возможности в программировании
        - Улучшенное следование инструкциям и генерация длинных текстов
        - Поддержка многоязычия (29+ языков)
        """)
    
    # Обработчики событий
    generate_btn.click(
        fn=generate_code_with_context,
        inputs=[prompt_input, file_input, max_length_slider, temperature_slider, top_p_slider],
        outputs=response_output
    )

# Запускаем приложение с MCP-сервером
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,  # Явно отключаем SSR
        mcp_server=True
    )