Update app.py

app.py

@@ -1,54 +1,27 @@
 import gradio as gr
-from ctransformers import AutoModelForCausalLM
-from huggingface_hub import hf_hub_download
+import requests
+import json
 from faster_whisper import WhisperModel
 
+OLLAMA_URL = "http://localhost:11434"
+
 MODELS = {
-    "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"
-    },
+    "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",
 }
 
-loaded_models = {}
-
 print("Loading Whisper...")
 whisper_model = WhisperModel("tiny", device="cpu", compute_type="int8")
 print("Whisper ready!")
 
-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 check_ollama():
+    try:
+        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
+        return r.status_code == 200
+    except:
+        return False
 
 def transcribe_audio(audio):
     if audio is None:
@@ -59,82 +32,110 @@ def transcribe_audio(audio):
     except Exception as e:
         return f"[STT Error: {e}]"
 
-def chat(message, history, model_name, temperature, max_tokens):
-    llm = get_model(model_name)
-    if llm is None:
-        return "❌ Model not found"
+def chat_stream(message, history, model_name, temperature, max_tokens):
+    if not check_ollama():
+        yield "⏳ Ollama starting... wait 30 seconds and try again."
+        return
     
-    prompt = "<|im_start|>system\nYou are an expert coding assistant. Always use markdown code blocks.<|im_end|>\n"
+    model = MODELS.get(model_name, "qwen2.5-coder:3b")
+    messages = [{"role": "system", "content": "You are an expert coding assistant. Always use markdown code blocks."}]
     
     for user_msg, assistant_msg in history:
-        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
+        messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
-            prompt += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
+            messages.append({"role": "assistant", "content": assistant_msg})
     
-    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+    messages.append({"role": "user", "content": message})
     
     try:
-        response = llm(prompt, max_new_tokens=max_tokens, temperature=temperature)
-        return response
+        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
     except Exception as e:
-        return f"Error: {e}"
+        yield 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..."
     
-    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"
+    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."
     
     try:
-        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
+        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}"
     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..."
     
-    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"
+    model = MODELS.get(model_name, "qwen2.5-coder:3b")
     
     try:
-        return llm(prompt, max_new_tokens=max_tokens, temperature=0.5)
+        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}"
     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..."
     
-    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"
+    model = MODELS.get(model_name, "qwen2.5-coder:3b")
+    prompt = f"Fix this code:\n```\n{code}\n```\nError: {error or 'Not working'}"
     
     try:
-        return llm(prompt, max_new_tokens=max_tokens, temperature=0.3)
+        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}"
     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**CTransformers Edition** • Any GGUF • No rate limits!")
+    gr.Markdown("# 🔥 Axon v5.1\n**Ollama Edition** • Qwen2.5-Coder + Qwen3-Coder • No rate limits!")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="Qwen2.5-Coder 3B (Fast)", label="🤖 Model")
@@ -177,9 +178,8 @@ 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 []
-        response = chat(message, history, model, temp, tokens)
-        history.append([message, response])
-        return history, ""
+        for chunk in chat_stream(message, history, model, temp, tokens):
+            yield history + [[message, chunk]], ""
     
     msg.submit(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
     send.click(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
@@ -189,7 +189,4 @@ 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)