app.py

import gradio as gr
from llama_cpp import Llama
import os

# Загрузка модели (GGUF из HF)
model_path = "eliza555beth2002/DeepSeek-R1-Distill-Text2SQL-OneEpoch-GGUF-q4"
llm = Llama.from_pretrained(
    model_path,
    n_ctx=2048,  # Контекст для схемы БД
    n_gpu_layers=-1,  # Использовать GPU, если доступно
    verbose=False
)

def generate_sql(natural_query, db_schema):
    # Шаблон промпта для Text2SQL (адаптируйте, если у модели другой)
    prompt = f"""You are a SQL expert. Given the database schema below, generate a valid SQL query for the natural language question.

    Database Schema:
    {db_schema}

    Question: {natural_query}

    SQL Query:"""
    
    output = llm(
        prompt,
        max_tokens=150,
        stop=[";", "\n\n"],
        echo=False
    )
    sql_query = output['choices'][0]['text'].strip()
    return sql_query

# Gradio интерфейс
with gr.Blocks(title="Text2SQL Demo") as demo:
    gr.Markdown("# Text2SQL: Преобразование текста в SQL")
    with gr.Row():
        natural_input = gr.Textbox(label="Вопрос на естественном языке", placeholder="Найди всех пользователей старше 30 лет")
        schema_input = gr.Textbox(label="Схема БД (таблицы и поля)", placeholder="CREATE TABLE users (id INT, name VARCHAR, age INT);")
    output = gr.Textbox(label="Сгенерированный SQL")
    submit_btn = gr.Button("Генерировать SQL")
    
    submit_btn.click(
        fn=generate_sql,
        inputs=[natural_input, schema_input],
        outputs=output
    )
    gr.Examples(
        examples=[
            ["Сколько продуктов дороже 100?", "CREATE TABLE products (id INT, name VARCHAR, price DECIMAL);"],
            ["Пользователи из Москвы", "CREATE TABLE users (id INT, name VARCHAR, city VARCHAR);"]
        ],
        inputs=[natural_input, schema_input]
    )

if __name__ == "__main__":
    demo.launch()