app.py

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Используем модель GPT-2 для генерации текста
model_name = "gpt2"  # Модель GPT-2
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def respond(message, history=None, system_message=None, max_tokens=512, temperature=0.7, top_p=0.95):
    if history is None:
        history = []  # Инициализируем пустой список, если history не передан

    # Объединяем сообщения в историю
    input_text = "\n".join([msg[1] for msg in history] + [message])

    # Токенизация текста
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)

    # Генерация ответа
    outputs = model.generate(
        inputs["input_ids"],
        max_length=max_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Формируем ответ согласно шаблону
    response = format_response(response)

    # Обновляем историю
    history.append((message, response))

    return response, history

def format_response(response):
    # Форматируем ответ в соответствии с шаблоном
    diagnosis = extract_diagnosis(response)
    operation = extract_operation(response)
    treatment = extract_treatment(response)

    formatted_response = f"Предварительный диагноз: {diagnosis}\nОперация: {operation}\nЛечение: {treatment}"

    return formatted_response

def extract_diagnosis(response):
    # Простой способ извлечь диагноз (можно улучшить с помощью NLP)
    diagnosis = response.split(".")[0]  # Пример: диагноз - первая часть ответа
    return diagnosis.strip()

def extract_operation(response):
    # Извлекаем название операции из ответа
    operation = "Не требуется"  # Пример, что операция не требуется
    return operation.strip()

def extract_treatment(response):
    # Извлекаем лечение (например, лечение как последняя часть ответа)
    treatment = response.split(".")[-1]  # Пример: лечение - последняя часть
    return treatment.strip()

# Интерфейс Gradio
demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox(value="Здравствуйте. Отвечай кратко...", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p"),
    ],
    outputs=["text", "state"],  # Возвращаем ответ и обновленную историю
    live=True  # Обновляем интерфейс в реальном времени
)

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