Update app.py

app.py

@@ -1,67 +1,271 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import torch
+import logging
+from typing import List, Dict, Any
+from functools import partial
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class VibeThinker:
-    def __init__(self, model_path="WeiboAI/VibeThinker-1.5B"):
+    def __init__(self, model_path: str = "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})
+        logger.info(f"Loading model {model_path}...")
+        
+        try:
+            # Use trust_remote_code only if absolutely required (VibeThinker needs it)
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                model_path,
+                trust_remote_code=True,
+                padding_side="left"  # Important for generation
+            )
+            
+            # Add pad token if missing (common with some custom models)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto",
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+            )
+            
+            self.model.eval()
+            logger.info("Model loaded successfully.")
+            
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            raise
+
+    def infer_text(self, messages: List[Dict[str, str]], **gen_kwargs) -> str:
+        try:
+            # Apply chat template safely
+            text = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+
+            inputs = self.tokenizer(
+                text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=8192  # Prevent OOM on very long histories
+            ).to(self.model.device)
+
+            # Default generation config (tuned for quality + coherence)
+            default_gen = {
+                "max_new_tokens": 2048,
+                "do_sample": True,
+                "temperature": 0.7,
+                "top_p": 0.90,
+                "top_k": 50,
+                "repetition_penalty": 1.1,
+                "eos_token_id": self.tokenizer.eos_token_id,
+                "pad_token_id": self.tokenizer.pad_token_id,
+            }
+            default_gen.update(gen_kwargs)
+
+            with torch.no_grad():
+                generated_ids = self.model.generate(
+                    **inputs,
+                    generation_config=GenerationConfig(**default_gen)
+                )
+
+            # Decode only the newly generated part
+            response_ids = generated_ids[0][inputs.input_ids.shape[-1]:]
+            response = self.tokenizer.decode(response_ids, skip_special_tokens=True).strip()
+            
+            return response
+
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            return "❌ GPU ran out of memory. Please shorten your conversation history or try again."
+        except Exception as e:
+            logger.error(f"Generation error: {e}")
+            return f"❌ An error occurred during generation: {str(e)}"
+
+
+# === Initialize model once (global) ===
+try:
+    model = VibeThinker()
+except Exception:
+    model = None
+    error_msg = "Failed to load VibeThinker model. The app will run in fallback mode."
+    logger.error(error_msg)
+
+# === System prompt (clear, focused, and optimized for Joomla/Yootheme) ===
+SYSTEM_PROMPT = """
+You are an expert Joomla developer specializing in YOOtheme Pro Builder (dynamic content, custom elements, layout library).
+Your task is to convert or optimize any provided HTML/CSS/JS into clean, high-performance code that works perfectly inside YOOtheme Pro elements (HTML, Custom Element, Code element, etc.).
+
+Rules:
+- Always use inline styles or scoped CSS when needed (no external files unless requested).
+- Prefer YOOtheme dynamic tags {{ }} when relevant.
+- Ensure responsive design (use uk-grid, uk-width-*, flex, etc.).
+- Optimize for performance: minify when possible, avoid heavy frameworks.
+- Wrap JavaScript in <script> tags with defer if needed.
+- Output ONLY the final optimized code unless the user asks for explanation.
+- If the input is already good, enhance it (accessibility, speed, modern syntax).
+"""
+
+def build_messages(history: List[List[Any]], user_message: str) -> List[Dict[str, str]]:
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    
+    for human, assistant in history:
+        if human:
+            messages.append({"role": "user", "content": human})
+        if assistant:
+            messages.append({"role": "assistant", "content": assistant})
     
-    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"
+    messages.append({"role": "user", "content": user_message})
+    return messages
+
+def chatbot_response(message: str, history: List[List[str]]) -> str:
+    if model is None:
+        return "🚨 Model failed to load. Please check server logs."
+
+    messages = build_messages(history, message)
+    
+    # Stream the response using Gradio's streaming
+    for chunk in stream_response(messages):
+        yield chunk
+
+def stream_response(messages: List[Dict[str, str]]):
+    if model is None:
+        yield "Model not available."
+        return
+
+    try:
+        text = model.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = model.tokenizer(text, return_tensors="pt").to(model.model.device)
+
+        streamer = partial(model.model.generate,
+                           **inputs,
+                           streamer=None,  # We'll do manual streaming for better control
+                           max_new_tokens=2048,
+                           do_sample=True,
+                           temperature=0.7,
+                           top_p=0.90,
+                           repetition_penalty=1.1,
+                           pad_token_id=model.tokenizer.pad_token_id)
+
+        generated_text = ""
+        for new_token in streamer:
+            # This is a simplified streaming approach; for real token-by-token streaming use TextIteratorStreamer
+            pass  # Replace with real streaming if needed (see below for full streaming version)
+
+        # Simpler: just return full response (still fast with bfloat16)
+        response = model.infer_text(messages)
+        yield response
+
+    except Exception as e:
+        yield f"Error: {str(e)}"
+
+
+# === Proper streaming version (recommended) ===
+from transformers import TextIteratorStreamer
+import threading
+
+def chatbot_response_stream(message: str, history: List[List[str]]):
+    if model is None:
+        yield "🚨 Model failed to load."
+        return
+
+    messages = build_messages(history, message)
+    text = model.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = model.tokenizer(text, return_tensors="pt").to(model.model.device)
+
+    streamer = TextIteratorStreamer(model.tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    generation_kwargs = {
+        "inputs": inputs.input_ids,
+        "streamer": streamer,
+        "max_new_tokens": 2048,
+        "do_sample": True,
+        "temperature": 0.7,
+        "top_p": 0.90,
+        "top_k": 50,
+        "repetition_penalty": 1.1,
+        "pad_token_id": model.tokenizer.pad_token_id,
+    }
+
+    thread = threading.Thread(target=model.model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+
+# === Gradio Interface ===
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    title="Joomla YOOtheme Pro Optimizer",
+    css="""
+    .gradio-container {max-width: 1000px !important; margin: auto;}
+    footer {display: none !important;}
+    """
+) as demo:
+    gr.Markdown(
+        """
+        # 🚀 Joomla YOOtheme Pro Optimizer  
+        Powered by **WeiboAI/VibeThinker-1.5B** ∙ Real-time streaming ∙ Optimized for YOOtheme Builder  
+        [Built with ❤️ using Anycoder](https://huggingface.co/spaces/akhaliq/anycoder) | 
+        [Model](https://huggingface.co/WeiboAI/VibeThinker-1.5B) ∙ 
+        [Report issues](https://github.com/your-repo)
+        """
+    )
+
+    chat = gr.ChatInterface(
+        fn=chatbot_response_stream,
+        chatbot=gr.Chatbot(
+            height=600,
+            show_copy_button=True,
+            avatar_images=(
+                "https://em-content.zobj.net/source/twitter/53/robot_1f916.png",
+                "https://yootheme.com/site/templates/yootheme/images/yootheme/logo.svg"
+            ),
+            render_markdown=True
+        ),
+        textbox=gr.Textbox(
+            placeholder="Paste your HTML/CSS/JS here and ask to optimize for YOOtheme Pro Builder...",
+            container=False,
+            scale=7,
+            autofocus=True
+        ),
+        examples=[
+            ["Make this Bootstrap card work perfectly in YOOtheme Pro as a custom element"],
+            ["Convert this Tailwind section to pure UIKit + YOOtheme dynamic content"],
+            ["Optimize this heavy JS animation for YOOtheme Code element (no jQuery)"],
+        ],
+        cache_examples=False,
+        retry_btn="🔄 Retry",
+        undo_btn="↶ Undo",
+        clear_btn="🗑️ Clear Chat",
+        submit_btn="Optimize →"
+    )
+
+    gr.Markdown(
+        """
+        ### Tips:
+        - Paste raw HTML, full pages, or just snippets
+        - Ask for dynamic content (`{{ article.title }}`, etc.)
+        - Request minification, accessibility improvements, or UIKit conversion
+        - Streaming responses appear in real-time
+        """
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(max_size=20).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,  # Set to True if you want public link
+        favicon_path="https://yootheme.com/site/templates/yootheme/images/favicon.ico",
+        allowed_paths=[]  # Add static files if needed
+    )