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

# Load model and tokenizer locally
MODEL_NAME = "BICORP/Lake-1-12B-spe"  # Replace with your local model directory if different

def load_model():
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
        return tokenizer, model
    except Exception as e:
        raise RuntimeError(f"Error loading model: {str(e)}")

tokenizer, model = load_model()

def generate_text(prompt, max_length=100, temperature=0.9):
    try:
        # Encode the input prompt
        inputs = tokenizer.encode(prompt, return_tensors='pt')
        
        # Generate output
        outputs = model.generate(
            inputs,
            max_length=max_length,
            temperature=temperature,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True
        )
        
        # Decode and return the generated text
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return generated_text
    except Exception as e:
        return f"Error generating text: {str(e)}"

# Create Gradio interface
interface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Input Prompt", lines=3),
        gr.Slider(50, 500, value=100, label="Max Length"),
        gr.Slider(0.1, 2.0, value=0.9, label="Temperature")
    ],
    outputs=gr.Textbox(label="Generated Text", lines=5),
    title="Local Model Demo - Text Generation",
    description="A locally loaded GPT-2 model for text generation",
    examples=[
        ["Once upon a time, in a land far away,"],
        ["The future of artificial intelligence", 150, 0.7],
        ["In a world where robots rule,", 200, 1.2]
    ]
)

# Run the app
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)