app.py

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

MODEL = "nn-tech/MetalGPT-1"

tokenizer = AutoTokenizer.from_pretrained(MODEL)

model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.float16,
    device_map="auto"
)

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

    output = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7
    )

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

demo = gr.Interface(
    fn=generate,
    inputs="text",
    outputs="text",
    title="MetalGPT Demo"
)

demo.launch()