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

MODEL_PATH = "rumeshprasanga6/PromptProAI"

model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
tokenizer.pad_token = tokenizer.eos_token

def respond(message, history):
    chat = f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    inputs = tokenizer(chat, return_tensors="pt").to("cuda")
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=256, temperature=0.7, top_p=0.9)
    result = tokenizer.decode(out[0], skip_special_tokens=True)
    answer = result.split("<|im_end|>")[0].split("<|im_start|>assistant\n")[-1].strip()
    return answer

demo = gr.ChatInterface(respond, title="PromptPro AI", description="Your prompt engineering AI assistant!")
demo.launch()