| 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() |