app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# Загружаем базовую модель и токенизатор
base_model_name = "t-tech/T-lite-it-1.0"
lora_repo = "shao3d/my-t-lite-qlora"  # Замени на твой логин

tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    device_map="cpu",  # Используем CPU для бесплатного Space
    torch_dtype=torch.float16  # FP16 для экономии памяти
)
model = PeftModel.from_pretrained(base_model, lora_repo)
model.eval()  # Переключаем в режим предсказания

def generate_response(history):
    # Если history пустой, берем сообщение из текстового поля (msg)
    if not history:
        # Получаем сообщение из текстового поля через Gradio (предполагаем, что оно передается отдельно)
        raise gr.Error("Пожалуйста, введите сообщение в текстовое поле перед отправкой.")  # Проверка на случай, если сообщение не введено
    # Берем последний вопрос пользователя
    user_message = history[-1][0] if history else gr.State.value  # Если history пуст, берем сообщение из State (опционально)
    # Генерируем ответ
    inputs = tokenizer(user_message, return_tensors="pt").to("cpu")
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,  # Уменьшили для скорости
        temperature=0.1,    # Сделали ответы точнее и быстрее
        top_p=0.5,          # Уменьшили для скорости
        do_sample=True      # Оставляем для разнообразия
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Возвращаем обновлённую историю: добавляем пару [вопрос, ответ]
    return history + [[user_message, response[len(user_message):].strip()]]

# Создаём интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Тест дообученной T-Lite")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Напиши сообщение для модели...")
    clear = gr.Button("Очистить чат")
    msg.submit(generate_response, inputs=[chatbot], outputs=chatbot)  # Указываем, что входные данные — это chatbot
    clear.click(lambda: [], None, chatbot)

demo.launch()