Update app.py

app.py

@@ -1,35 +1,44 @@
 import spaces
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import Qwen2ForCausalLM, Qwen2Tokenizer, Qwen2Config
 
-# 1. Configuration & Model Loading
+# 1. Configuration
 model_id = "erdemozkan/YOLO-7B-Qwen-Coder"
 
-print(f"🚀 Initializing YOLO CODER: {model_id}")
+print(f"🚀 YOLO CODER: Forcing Qwen2 Engine for {model_id}...")
 
-# FIX: Set use_fast=False and trust_remote_code=True
-tokenizer = AutoTokenizer.from_pretrained(
-    model_id, 
-    trust_remote_code=True,
-    use_fast=False # <--- THIS BYPASSES THE ERROR
-)
+# 2. Direct Loading (Bypassing AutoModel/AutoTokenizer)
+# We use the specific Qwen2 classes to avoid the 'Unrecognized' error
+try:
+    tokenizer = Qwen2Tokenizer.from_pretrained(
+        model_id, 
+        trust_remote_code=True,
+        use_fast=False # Standard Python backend is more stable for custom weights
+    )
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    torch_dtype=torch.bfloat16, 
-    device_map="auto",
-    trust_remote_code=True
-)
+    model = Qwen2ForCausalLM.from_pretrained(
+        model_id, 
+        torch_dtype=torch.bfloat16, 
+        device_map="auto",
+        trust_remote_code=True
+    )
+except Exception as e:
+    print(f"❌ Direct load failed: {e}")
+    print("Falling back to Auto-classes with forced config...")
+    # Last ditch effort if Qwen2 classes aren't in the path
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, device_map="auto")
 
-# 2. The Core Logic
+# 3. The Core Logic
 @spaces.GPU(duration=60) 
 def yoco_heal(broken_code):
     if not broken_code or not broken_code.strip():
-        return "⚠️ YOLO needs code. Paste something!"
+        return "⚠️ YOLO needs code."
 
     prompt = (
-        "### System: You are YOLO CODER. An elite autonomous agent.\n"
+        "### System: You are YOLO CODER (yoco). An elite autonomous agent.\n"
         f"### Broken Code:\n{broken_code}\n"
         "### Fixed Code (No talk, just code):"
     )
@@ -42,34 +51,28 @@ def yoco_heal(broken_code):
             max_new_tokens=1024, 
             temperature=0.2, 
             do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id else 151643
         )
     
     full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     marker = "### Fixed Code (No talk, just code):"
     return full_response.split(marker)[1].strip() if marker in full_response else full_response
 
-# 3. UI and CSS
-css_hack = """
+# 4. Neon-Hacker UI
+css = """
 .gradio-container { background-color: #000 !important; }
-#yoco-header h1 { color: #ffff00 !important; text-shadow: 0 0 10px #ffff00; text-align: center; }
-.code-container { border: 1px solid #ffff00 !important; background: #09090b !important; }
-#yolo-btn { background: #ffff00 !important; color: #000 !important; font-weight: 900 !important; box-shadow: 0 0 15px #ffff00; }
-footer { display: none !important; }
+#header h1 { color: #ffff00 !important; text-shadow: 0 0 10px #ffff00; text-align: center; }
+.code-box { border: 1px solid #ffff00 !important; background: #09090b !important; }
+#yolo-btn { background: #ffff00 !important; color: #000 !important; font-weight: 900 !important; }
 """
 
-with gr.Blocks(css=css_hack) as demo:
-    with gr.Column(elem_id="yoco-header"):
-        gr.Markdown("# YOLO CODER")
-    
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# YOLO CODER", elem_id="header")
     with gr.Row():
-        with gr.Column():
-            input_box = gr.Code(label="INPUT", language="python", lines=15, elem_classes="code-container")
-            heal_btn = gr.Button("YOLO IT!", elem_id="yolo-btn")
-        with gr.Column():
-            output_box = gr.Code(label="HEALED", language="python", lines=15, elem_classes="code-container", interactive=False)
-
-    heal_btn.click(fn=yoco_heal, inputs=input_box, outputs=output_box)
+        in_code = gr.Code(label="INPUT", language="python", lines=15, elem_classes="code-box")
+        out_code = gr.Code(label="HEALED", language="python", lines=15, elem_classes="code-box", interactive=False)
+    btn = gr.Button("YOLO IT!", elem_id="yolo-btn")
+    btn.click(yoco_heal, in_code, out_code)
 
 if __name__ == "__main__":
     demo.launch()