Update app.py

changed model to llama 7b

app.py

@@ -1,26 +1,29 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from llama_cpp import Llama
 
-# Load the small model (WizardCoder-1B)
-model_name = "WizardLM/WizardCoder-1B-V1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")
+# Path to model (downloaded automatically from HF hub)
+MODEL_PATH = "TheBloke/CodeLlama-7B-Instruct-GGUF"
 
-def generate_manim_code(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.7)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+# Load LLM (first time takes time)
+llm = Llama.from_pretrained(
+    repo_id=MODEL_PATH,
+    filename="codellama-7b-instruct.Q4_K_M.gguf",  # 4-bit quantized version
+    n_ctx=2048,
+    verbose=True
+)
 
-description = "💡 Enter a prompt like: *Write a Manim script to animate a growing square*"
+# Prompt wrapper
+def build_prompt(user_prompt):
+    return f"[INST] Write Python code using the Manim library: {user_prompt} [/INST]"
 
-iface = gr.Interface(
-    fn=generate_manim_code,
-    inputs=gr.Textbox(label="Prompt", placeholder="e.g., Create a bouncing ball using Manim"),
-    outputs=gr.Textbox(label="Generated Manim Code"),
-    title="🎬 Manim Code Generator",
-    description=description
-)
+def generate_code(prompt):
+    result = llm(build_prompt(prompt), max_tokens=512, temperature=0.7, stop=["</s>"])
+    return result["choices"][0]["text"]
 
-iface.launch()
+gr.Interface(
+    fn=generate_code,
+    inputs=gr.Textbox(label="Prompt", placeholder="e.g., Animate a bouncing ball in Manim"),
+    outputs=gr.Textbox(label="Generated Python Code"),
+    title="🐍 Manim Code Generator - Code Llama 7B",
+    description="Powered by llama-cpp and Code Llama 7B (Quantized). Runs on CPU!"
+).launch()