Update app.py

app.py

@@ -1,143 +1,140 @@
 import gradio as gr
-import requests
-import json
+from ctransformers import AutoModelForCausalLM
+from huggingface_hub import hf_hub_download
 from faster_whisper import WhisperModel
 
-OLLAMA_URL = "http://localhost:11434"
-
 MODELS = {
-    "Qwen2.5-Coder 1.5B (Fastest)": "qwen2.5-coder:1.5b",
-    "Qwen2.5-Coder 3B (Fast)": "qwen2.5-coder:3b",
-    "Qwen2.5-Coder 7B (Quality)": "qwen2.5-coder:7b",
-    "Qwen3-Coder 30B-A3B (Best)": "qwen3-coder:30b-a3b",
+    "Qwen2.5-Coder 3B (Fast)": {
+        "repo": "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
+        "file": "qwen2.5-coder-3b-instruct-q4_k_m.gguf",
+        "type": "qwen2"
+    },
+    "Qwen2.5-Coder 7B (Quality)": {
+        "repo": "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
+        "file": "qwen2.5-coder-7b-instruct-q4_k_m.gguf",
+        "type": "qwen2"
+    },
+    "Qwen3-Coder 30B-A3B (Best)": {
+        "repo": "Qwen/Qwen3-Coder-30B-A3B-Instruct-GGUF",
+        "file": "qwen3-coder-30b-a3b-instruct-q4_k_m.gguf",
+        "type": "qwen2"
+    },
 }
 
+loaded_models = {}
+
 print("Loading Whisper...")
 whisper_model = WhisperModel("tiny", device="cpu", compute_type="int8")
 print("Whisper ready!")
 
-def check_ollama():
-    try:
-        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
-        return r.status_code == 200
-    except:
-        return False
+def get_model(model_name):
+    if model_name in loaded_models:
+        return loaded_models[model_name]
+    
+    info = MODELS.get(model_name)
+    if not info:
+        return None
+    
+    print(f"Downloading {model_name}...")
+    path = hf_hub_download(repo_id=info["repo"], filename=info["file"])
+    
+    print(f"Loading {model_name}...")
+    llm = AutoModelForCausalLM.from_pretrained(
+        path,
+        model_type=info["type"],
+        context_length=4096,
+        threads=4
+    )
+    loaded_models[model_name] = llm
+    print(f"{model_name} ready!")
+    
+    return llm
 
 def transcribe_audio(audio):
     if audio is None:
         return ""
-    
     try:
         segments, _ = whisper_model.transcribe(audio)
-        text = " ".join([seg.text for seg in segments])
-        return text.strip()
+        return " ".join([seg.text for seg in segments]).strip()
     except Exception as e:
         return f"[STT Error: {e}]"
 
-def chat_stream(message, history, model_name, temperature, max_tokens):
-    if not check_ollama():
-        yield "⏳ Ollama starting... wait 30 seconds and try again."
-        return
+def chat(message, history, model_name, temperature, max_tokens):
+    llm = get_model(model_name)
+    if llm is None:
+        return "❌ Model not found"
     
-    model = MODELS.get(model_name, "qwen2.5-coder:3b")
-    messages = [{"role": "system", "content": "You are an expert coding assistant. Always use markdown code blocks."}]
+    prompt = "<|im_start|>system\nYou are an expert coding assistant. Always use markdown code blocks.<|im_end|>\n"
     
     for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
+        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
         if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
+            prompt += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
     
-    messages.append({"role": "user", "content": message})
+    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
     
     try:
-        response = requests.post(
-            f"{OLLAMA_URL}/api/chat",
-            json={"model": model, "messages": messages, "stream": True, "options": {"temperature": temperature, "num_predict": max_tokens}},
-            stream=True, timeout=300
-        )
-        
-        full = ""
-        for line in response.iter_lines():
-            if line:
-                try:
-                    data = json.loads(line)
-                    if "message" in data:
-                        full += data["message"].get("content", "")
-                        yield full
-                except:
-                    continue
+        response = llm(prompt, max_new_tokens=max_tokens, temperature=temperature)
+        return response
     except Exception as e:
-        yield f"Error: {e}"
+        return f"Error: {e}"
 
 def generate_code(prompt, language, model_name, max_tokens):
     if not prompt.strip():
         return "Please describe what you want."
-    if not check_ollama():
-        return "⏳ Ollama starting..."
     
-    model = MODELS.get(model_name, "qwen2.5-coder:3b")
-    full_prompt = f"Write {language} code for: {prompt}\n\nOutput ONLY code in a markdown block."
+    llm = get_model(model_name)
+    if llm is None:
+        return "❌ Model not found"
+    
+    full_prompt = f"<|im_start|>user\nWrite {language} code for: {prompt}\n\nOutput ONLY code in a markdown block.<|im_end|>\n<|im_start|>assistant\n"
     
     try:
-        r = requests.post(
-            f"{OLLAMA_URL}/api/generate",
-            json={"model": model, "prompt": full_prompt, "stream": False, "options": {"temperature": 0.3, "num_predict": max_tokens}},
-            timeout=300
-        )
-        if r.status_code == 200:
-            result = r.json().get("response", "")
-            if "```" in result:
-                parts = result.split("```")
-                if len(parts) >= 2:
-                    code = parts[1]
-                    if "\n" in code:
-                        code = code.split("\n", 1)[-1]
-                    return code.strip()
-            return result
-        return f"Error: {r.text}"
+        result = llm(full_prompt, max_new_tokens=max_tokens, temperature=0.3)
+        if "```" in result:
+            parts = result.split("```")
+            if len(parts) >= 2:
+                code = parts[1]
+                if "\n" in code:
+                    code = code.split("\n", 1)[-1]
+                return code.strip()
+        return result
     except Exception as e:
         return f"Error: {e}"
 
 def explain_code(code, model_name, max_tokens):
     if not code.strip():
         return "Paste code to explain."
-    if not check_ollama():
-        return "⏳ Ollama starting..."
     
-    model = MODELS.get(model_name, "qwen2.5-coder:3b")
+    llm = get_model(model_name)
+    if llm is None:
+        return "❌ Model not found"
+    
+    prompt = f"<|im_start|>user\nExplain this code:\n```\n{code}\n```<|im_end|>\n<|im_start|>assistant\n"
     
     try:
-        r = requests.post(
-            f"{OLLAMA_URL}/api/generate",
-            json={"model": model, "prompt": f"Explain this code:\n```\n{code}\n```", "stream": False, "options": {"num_predict": max_tokens}},
-            timeout=300
-        )
-        return r.json().get("response", "") if r.status_code == 200 else f"Error: {r.text}"
+        return llm(prompt, max_new_tokens=max_tokens, temperature=0.5)
     except Exception as e:
         return f"Error: {e}"
 
 def fix_code(code, error, model_name, max_tokens):
     if not code.strip():
         return "Paste code to fix."
-    if not check_ollama():
-        return "⏳ Ollama starting..."
     
-    model = MODELS.get(model_name, "qwen2.5-coder:3b")
-    prompt = f"Fix this code:\n```\n{code}\n```\nError: {error or 'Not working'}"
+    llm = get_model(model_name)
+    if llm is None:
+        return "❌ Model not found"
+    
+    prompt = f"<|im_start|>user\nFix this code:\n```\n{code}\n```\nError: {error or 'Not working'}<|im_end|>\n<|im_start|>assistant\n"
     
     try:
-        r = requests.post(
-            f"{OLLAMA_URL}/api/generate",
-            json={"model": model, "prompt": prompt, "stream": False, "options": {"temperature": 0.3, "num_predict": max_tokens}},
-            timeout=300
-        )
-        return r.json().get("response", "") if r.status_code == 200 else f"Error: {r.text}"
+        return llm(prompt, max_new_tokens=max_tokens, temperature=0.3)
     except Exception as e:
         return f"Error: {e}"
 
 with gr.Blocks(title="Axon v5.1", theme=gr.themes.Soft(primary_hue="purple")) as demo:
     
-    gr.Markdown("# 🔥 Axon v5.1\n**Ollama Edition** • Qwen2.5-Coder running locally • No rate limits!")
+    gr.Markdown("# 🔥 Axon v5.1\n**CTransformers Edition** • Any GGUF • No rate limits!")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="Qwen2.5-Coder 3B (Fast)", label="🤖 Model")
@@ -180,8 +177,9 @@ with gr.Blocks(title="Axon v5.1", theme=gr.themes.Soft(primary_hue="purple")) as
     
     def respond(message, history, model, temp, tokens):
         history = history or []
-        for chunk in chat_stream(message, history, model, temp, tokens):
-            yield history + [[message, chunk]], ""
+        response = chat(message, history, model, temp, tokens)
+        history.append([message, response])
+        return history, ""
     
     msg.submit(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
     send.click(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
@@ -191,4 +189,7 @@ with gr.Blocks(title="Axon v5.1", theme=gr.themes.Soft(primary_hue="purple")) as
     explain_btn.click(explain_code, [explain_input, model_dropdown, max_tokens], explain_output)
     fix_btn.click(fix_code, [fix_input, fix_error, model_dropdown, max_tokens], fix_output)
 
+print("Pre-loading default model...")
+get_model("Qwen2.5-Coder 3B (Fast)")
+
 demo.launch(server_name="0.0.0.0", server_port=7860)