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

class TextGenerator:
    def __init__(self, model_name, device='cpu'):
        self.device = device
        self.load_model(model_name)

    def load_model(self, model_name):
        # Load model and tokenizer from Hugging Face
        print("Loading model and tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
        self.model.to(self.device)
        print("Model loaded successfully!")

    def generate_text(self, prompt, max_length=100, temperature=0.7, top_k=50, top_p=0.9):
        # Tokenize input
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
        
        # Generate text
        with torch.no_grad():
            output_ids = self.model.generate(
                input_ids,
                max_length=max_length,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id
            )
        
        # Decode output tokens
        generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
        return generated_text

def create_gradio_interface(model_name):
    generator = TextGenerator(model_name)

    def generate(prompt, max_length, temperature, top_k, top_p):
        try:
            return generator.generate_text(
                prompt=prompt,
                max_length=max_length,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p
            )
        except Exception as e:
            return f"Error: {str(e)}"

    # Define Gradio interface
    interface = gr.Interface(
        fn=generate,
        inputs=[
            gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
            gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Maximum Length"),
            gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
            gr.Slider(minimum=0, maximum=100, value=50, step=5, label="Top-k"),
            gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p"),
        ],
        outputs=gr.Textbox(label="Generated Text"),
        title="Reality123b/Xylaria-1.4-Senoa-Test",
        description="Generate text using the Reality123b/Xylaria-1.4-Senoa-Test model optimized for CPU usage.",
    )
    return interface

if __name__ == "__main__":
    # Use the model from Hugging Face
    model_name = "Reality123b/Xylaria-1.4-Senoa-Test"
    
    # Create and launch Gradio interface
    interface = create_gradio_interface(model_name)
    interface.launch(share=True)