app.py

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

MODEL_NAME = "RameshRathod/ego-45m-pretrained"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

print("Loading YOUR model from Hugging Face...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"
)

def generate_text(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=120,
            temperature=0.8,
            top_k=40,
            do_sample=True
        )

    return tokenizer.decode(output[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(label="Type your prompt"),
    outputs=gr.Textbox(label="Model response"),
    title="🧠 Ego 45M — Your Custom LLM",
    description="This Space runs YOUR trained model: RameshRathod/ego-45m-pretrained"
)

demo.launch()

# "making the final push "