app.py

import gradio as gr
import os
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Путь для хранения модели локально в Space
model_dir = "./models"
os.makedirs(model_dir, exist_ok=True)

# Название модели и репозитория
repo_id = "Mykes/medicus"
model_filename = "medicus-Q2_K.gguf"

# Загружаем модель, если она еще не загружена
model_path = os.path.join(model_dir, model_filename)
if not os.path.exists(model_path):
    print(f"Downloading model {model_filename} from {repo_id}...")
    model_path = hf_hub_download(repo_id=repo_id, filename=model_filename, cache_dir=model_dir)
    print(f"Model downloaded to {model_path}")

# Загружаем модель в память
print("Loading model into memory...")
llm = Llama(model_path=model_path, n_ctx=2048, n_threads=4, n_batch=512,)  # n_ctx - максимальная длина контекста

def respond(
    message,
    history: list[tuple[str, str]],
    # system_message,
    max_tokens,
    temperature,
    top_p,
):
    # Формируем историю сообщений в формате, подходящем для модели
    # messages = [{"role": "system", "content": system_message}]

    # for val in history:
    #     if val[0]:
    #         messages.append({"role": "user", "content": val[0]})
    #     if val[1]:
    #         messages.append({"role": "assistant", "content": val[1]})

    # messages.append({"role": "user", "content": message})

    # Формируем prompt для модели (в GGUF обычно используется специфический формат)
    prompt = ""
    # for msg in messages:
    #     if msg["role"] == "system":
    #         prompt += f"System: {msg['content']}\n"
    #     elif msg["role"] == "user":
    #         prompt += f"User: {msg['content']}\n"
    #     elif msg["role"] == "assistant":
    #         prompt += f"Assistant: {msg['content']}\n"
    for user_msg, model_msg in history:
        if user_msg:
            prompt += f"<start_of_turn>user\n{user_msg}\n"
        if model_msg:
            prompt += f"<start_of_turn>model\n{model_msg}\n"

    # Добавляем текущее сообщение пользователя
    prompt += f"<start_of_turn>user\n{message}\n"

    # Генерируем ответ с помощью модели
    response = llm(
        prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        # stop=["User:", "System:"],  # Останавливаем генерацию, если начинается новое сообщение
        stop=["<start_of_turn>"],
        stream=True
    )

    # Потоковая обработка ответа
    full_response = ""
    for chunk in response:
        token = chunk["choices"][0]["text"]
        full_response += token
        yield full_response

# Настройка интерфейса Gradio
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        # gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

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