import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

base_model_id = "google/gemma-2b-it"
adapter_model_id = "MatclassAI/Matclass-Gemma-2b-v1"

print("Loading base tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(adapter_model_id)

print("Loading base framework layers...")
model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True
)

print("Injecting custom fine-tuned adapters into core matrices...")
model.load_adapter(adapter_model_id)
print("Model infrastructure fully compiled!")

def generate_text(prompt, system_prompt, temperature, max_tokens):
    if system_prompt:
        formatted_prompt = f"<start_of_turn>user\n{system_prompt}\n\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
    else:
        formatted_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
        
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=True,
            temperature=temperature,
            top_k=50,
            top_p=0.95
        )
        
    generated = tokenizer.decode(outputs[0], skip_special_tokens=False)
    
    if formatted_prompt in generated:
        return generated.split("<start_of_turn>model\n")[-1].replace("<end_of_turn>", "").strip()
    return generated

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Matclass AI - Gemma 2B Adapter Workspace")
    gr.Markdown("Dynamic LoRA adapter execution layer tracking custom weights.")
    
    with gr.Row():
        with gr.Column(scale=2):
            system_input = gr.Textbox(
                label="System Instructions", 
                placeholder="Define assistant behavior or context here...", 
                lines=2
            )
            user_input = gr.Textbox(
                label="Query", 
                placeholder="Ask anything...", 
                lines=5
            )
            submit_btn = gr.Button("Generate Text", variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown("### Generation Hyperparameters")
            temp_slider = gr.Slider(
                minimum=0.1, 
                maximum=1.5, 
                value=0.7, 
                step=0.05, 
                label="Temperature"
            )
            tokens_slider = gr.Slider(
                minimum=64, 
                maximum=1024, 
                value=256, 
                step=64, 
                label="Max New Tokens"
            )
            
    with gr.Row():
        output_text = gr.Textbox(label="Model Output", lines=8, interactive=False)
        
    submit_btn.click(
        fn=generate_text,
        inputs=[user_input, system_input, temp_slider, tokens_slider],
        outputs=output_text
    )

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