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,
):
    # Инициализация history как пустого списка, если его нет
    if history is None:
        history = []
    elif not isinstance(history, list):
        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)

    return response

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",
)

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