Spaces:
Sleeping
Sleeping
| 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: <html>...</html>"], | |
| retry_btn="Retry", | |
| undo_btn="Undo", | |
| clear_btn="Clear" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |