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

def generate_text(model_id, prompt, temperature, top_k, top_p, max_tokens, repetition_penalty):
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
    model.eval()
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=True,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            pad_token_id=tokenizer.eos_token_id
        )
    return tokenizer.decode(output[0], skip_special_tokens=True)

gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Model ID", placeholder="Enter Model ID"),
        gr.Textbox(label="Prompt", placeholder="Type something here...", lines=4),
        gr.Slider(0.1, 1.5, value=1.0, step=0.1, label="Temperature"),
        gr.Slider(1, 100, value=50, step=1, label="Top-K"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P"),
        gr.Slider(10, 512, value=128, step=1, label="Max New Tokens"),
        gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Repetition Penalty")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="🧠 AlphaMindQ Fork — Custom Hugging Face Model",
    theme="default"
).launch()