File size: 3,260 Bytes
c3dfc09
dfa1a47
c3dfc09
fd212a6
9d20ec4
dfa1a47
 
c3dfc09
f29c29c
133324c
f29c29c
ca1d8ee
f29c29c
133324c
8d3c493
f29c29c
dfa1a47
 
f29c29c
dfa1a47
 
 
8d3c493
 
dfa1a47
 
 
 
133324c
f29c29c
133324c
 
f29c29c
 
 
 
e1f2405
133324c
 
 
 
 
 
 
 
 
 
 
fd212a6
133324c
 
 
 
 
fd212a6
133324c
 
 
 
 
 
 
ca1d8ee
dfa1a47
 
 
ca1d8ee
133324c
ca1d8ee
 
c3dfc09
f29c29c
 
c3dfc09
 
133324c
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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()