import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
import os
import gc

class EngineAutoreg:
    def __init__(self):
        print("=" * 50)
        print("Loading Arch-Router with LoRA...")
        print("=" * 50)
        
        repo_name = "MarkProMaster229/experimental_models"
        lora_subfolder = "loraForArchkit/loraForArch4"
        base_model_name = "katanemo/Arch-Router-1.5B"
        
        # Определяем устройство
        if torch.cuda.is_available():
            device_map = "auto"
            torch_dtype = torch.float16
            print("Using GPU")
        else:
            device_map = "cpu"
            torch_dtype = torch.float32
            print("Using CPU (slow)")
        
        # Токенизатор
        self.tokenizer = AutoTokenizer.from_pretrained(
            base_model_name,
            trust_remote_code=True
        )
        
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        
        # Базовая модель
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            device_map=device_map,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        
        # LoRA
        self.model = PeftModel.from_pretrained(
            base_model,
            repo_name,
            subfolder=lora_subfolder,
            token=os.environ.get("HF_TOKEN")
        )
        
        self.model.eval()
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        print("Model loaded!")
    
    def generate(self, prompt, max_new_tokens=100, temperature=0.1):
        formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
        
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=temperature > 0,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            )
        
        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Извлекаем только ответ
        if "<|im_start|>assistant" in generated_text:
            response = generated_text.split("<|im_start|>assistant")[-1].strip()
        else:
            response = generated_text.replace(formatted_prompt, "").strip()
        
        return response

# Глобальная модель
print("Initializing model...")
engine = EngineAutoreg()

def generate_response(prompt, max_tokens, temperature):
    if not prompt.strip():
        return "Please enter a prompt"
    
    try:
        response = engine.generate(
            prompt=prompt,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature)
        )
        return response
    except Exception as e:
        return f"Error: {str(e)}"

# Создаем Gradio интерфейс
with gr.Blocks(theme=gr.themes.Soft(), title="Arch-Router + LoRA") as demo:
    gr.Markdown("""
    # 🤖 Arch-Router with Custom LoRA
    Base model: `katanemo/Arch-Router-1.5B`  
    LoRA adapter: `MarkProMaster229/experimental_models/loraForArchkit/loraForArch4`
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Your Prompt",
                placeholder="Enter your prompt here...",
                lines=4
            )
            
            with gr.Row():
                max_tokens = gr.Slider(
                    minimum=10,
                    maximum=500,
                    value=100,
                    step=10,
                    label="Max New Tokens"
                )
                temperature = gr.Slider(
                    minimum=0,
                    maximum=2,
                    value=0.1,
                    step=0.1,
                    label="Temperature"
                )
            
            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            output = gr.Textbox(
                label="Generated Response",
                lines=10,
                interactive=False
            )
    
    # Примеры
    gr.Examples(
        examples=[
            ["What is MVC architecture?", 100, 0.1],
            ["Explain microservices architecture", 150, 0.1],
            ["What is the difference between monolithic and microservices?", 200, 0.1],
        ],
        inputs=[prompt_input, max_tokens, temperature]
    )
    
    generate_btn.click(
        fn=generate_response,
        inputs=[prompt_input, max_tokens, temperature],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)