import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig import torch class VibeThinker: def __init__(self, model_path="WeiboAI/VibeThinker-1.5B"): self.model_path = model_path self.model = AutoModelForCausalLM.from_pretrained( self.model_path, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, device_map="auto" ) self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) def infer_text(self, messages): text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) generation_config = dict( max_new_tokens=4096, do_sample=True, temperature=0.6, top_p=0.95, top_k=None ) generated_ids = self.model.generate( **model_inputs, generation_config=GenerationConfig(**generation_config) ) generated_ids = [ output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) ] response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] return response model = VibeThinker() def chatbot_response(message, history): messages = [{"role": "system", "content": "You are a helpful assistant that optimizes HTML for Joomla Yootheme Builder. Generate optimized scripts based on user-provided HTML."}] for user_msg, assistant_msg in history: messages.append({"role": "user", "content": user_msg}) if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg}) messages.append({"role": "user", "content": message}) response = model.infer_text(messages) return response with gr.Blocks() as demo: gr.Markdown("# Joomla Yootheme Builder Optimizer\n[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)") gr.ChatInterface( chatbot_response, chatbot=gr.Chatbot(height=500), textbox=gr.Textbox(placeholder="Provide HTML to optimize for Joomla Yootheme Builder", container=False, scale=7), title="Joomla Optimizer Chatbot", description="Chat with the VibeThinker model to optimize your HTML for Joomla Yootheme Builder.", theme="soft", examples=["Optimize this HTML for Joomla: ..."], retry_btn="Retry", undo_btn="Undo", clear_btn="Clear" ) if __name__ == "__main__": demo.launch()