app.py

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

# Load model and tokenizer from Hugging Face Hub
model_name = "itriedcoding/Sage"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def generate_text(prompt, max_length, temperature):
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_length=int(max_length),
            temperature=temperature,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Prompt", value="Hello", placeholder="Enter your prompt here"),
        gr.Slider(minimum=10, maximum=100, value=30, label="Max Length"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.8, label="Temperature")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="🤖 Sage Text Generator",
    description="Generate text using the Sage custom character-level language model.",
    examples=[
        ["Hello", 30, 0.8],
        ["The weather", 30, 0.8],
        ["Deep learning", 30, 0.8]
    ]
)

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