import gradio as gr
from transformers import pipeline
import torch
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model_name = "microsoft/Phi-3-mini-4k-instruct"
try:
    logger.info(f"Попытка загрузки модели {model_name}...")
    generator = pipeline(
        "text-generation",
        model=model_name,
        device=-1,
        framework="pt",
        max_length=150,
        truncation=True,
        model_kwargs={"torch_dtype": torch.float32}
    )
    logger.info("Модель успешно загружена.")
except Exception as e:
    logger.error(f"Ошибка загрузки модели: {e}")
    exit(1)

def respond(message, max_tokens=150, temperature=0.5, top_p=0.8):
    prompt = f"You are a medical chatbot. User says: '{message}'. Provide diagnosis and treatment in English."
    try:
        logger.info(f"Генерация ответа для: {message}")
        outputs = generator(
            prompt,
            max_length=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            num_return_sequences=1
        )
        response = outputs[0]["generated_text"].replace(prompt, "").strip()
        # Простой перевод (для примера)
        if "Diagnosis" in response:
            diagnosis = response.split("Diagnosis:")[1].split("Treatment:")[0].strip()
            treatment = response.split("Treatment:")[1].strip()
            response = f"Диагноз: {diagnosis}\nЛечение: {treatment}"
    except Exception as e:
        logger.error(f"Ошибка генерации ответа: {e}")
        return f"Ошибка генерации: {e}"
    return response

demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox(label="Ваше сообщение", placeholder="Опишите симптомы (например, 'Болит горло')..."),
        gr.Slider(minimum=50, maximum=300, value=150, step=10, label="Макс. токенов"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Температура"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.8, label="Top-p")
    ],
    outputs="text",
    title="Медицинский чат-бот на базе Phi-3-mini-4k-instruct",
    theme=gr.themes.Soft()
)

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