app.py

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import gradio as gr

# ------------------------------
# Model & Adapter Paths
# ------------------------------
BASE_MODEL = "deepseek-ai/deepseek-coder-6.7b-base"
ADAPTER_PATH = "Agasthya0/colabmind-coder-6.7b-ml-qlora"

# ------------------------------
# Load Tokenizer
# ------------------------------
print("🚀 Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

# ------------------------------
# Load Base Model in 4-bit (for Hugging Face Spaces GPU)
# ------------------------------
print("🧠 Loading base model in 4-bit...")
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    load_in_4bit=True,
    device_map="auto",
    torch_dtype=torch.float16
)

# ------------------------------
# Load LoRA Adapter
# ------------------------------
print("🔗 Attaching LoRA adapter...")
model = PeftModel.from_pretrained(model, ADAPTER_PATH)

# ------------------------------
# Inference Function
# ------------------------------
def generate_code(prompt):
    """Generate Python/ML code based on the user prompt."""
    if not prompt.strip():
        return "⚠️ Please enter a prompt."

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.2,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# ------------------------------
# Custom HTML Frontend Loader
# ------------------------------
def load_html():
    """Loads your custom HTML page."""
    with open("index.html", "r", encoding="utf-8") as f:
        return f.read()

# ------------------------------
# Gradio App (Custom HTML + API)
# ------------------------------
with gr.Blocks(css=None) as demo:
    # Load your custom webpage
    gr.HTML(load_html())

    # Hidden API interface for JS fetch requests
    with gr.Row(visible=False):
        prompt_box = gr.Textbox()
        output_box = gr.Textbox()

    # API endpoint for your webpage's JS to call
    prompt_box.change(
        fn=generate_code,
        inputs=prompt_box,
        outputs=output_box
    )

# ------------------------------
# Launch App
# ------------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)