codex_cli.py

#!/usr/bin/env python3
"""
Codex-like Open-Source Code Generator CLI
Generates code from natural language prompts using CodeLLaMA
"""

import click
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from pathlib import Path

# Configuration
MODEL_NAME = "meta-llama/CodeLlama-7b-hf"  # or "bigcode/starcoder" as alternative
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

class CodeGenerator:
    def __init__(self, model_name: str = MODEL_NAME):
        self.model_name = model_name
        self.device = DEVICE
        self.tokenizer = None
        self.model = None
        self.load_model()

    def load_model(self):
        """Load the model and tokenizer"""
        click.echo(f"Loading model: {self.model_name} on {self.device.upper()}...")
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                device_map="auto",
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
            )
            click.echo("✓ Model loaded successfully!")
        except Exception as e:
            click.echo(f"✗ Error loading model: {e}", err=True)
            click.echo("Make sure you have HuggingFace credentials set up for CodeLlama access", err=True)
            raise

    def generate(self, prompt: str, max_length: int = 256, temperature: float = 0.7) -> str:
        """Generate code from a prompt"""
        if not self.model:
            raise RuntimeError("Model not loaded")
        
        # Format the prompt
        formatted_prompt = f"<s>[INST] {prompt} [/INST]"
        
        # Tokenize
        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
        
        # Generate
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=temperature,
                top_p=0.95,
                do_sample=True,
                eos_token_id=self.tokenizer.eos_token_id
            )
        
        # Decode
        generated_code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Clean up the output
        if "[/INST]" in generated_code:
            generated_code = generated_code.split("[/INST]")[1].strip()
        
        return generated_code

@click.command()
@click.option(
    '--prompt', '-p',
    type=str,
    required=True,
    help='Natural language description of the code to generate'
)
@click.option(
    '--language', '-l',
    type=str,
    default='python',
    help='Programming language (python, javascript, go, etc.)'
)
@click.option(
    '--length',
    type=int,
    default=256,
    help='Maximum length of generated code'
)
@click.option(
    '--temperature',
    type=float,
    default=0.7,
    help='Creativity level (0.0 = deterministic, 1.0 = random)'
)
@click.option(
    '--output', '-o',
    type=click.Path(),
    help='Save output to file'
)
def main(prompt: str, language: str, length: int, temperature: float, output: str):
    """
    AI-powered code generator using CodeLlama.
    
    Example:
        codex-cli -p "Write a function to check if a number is prime" -l python
    """
    try:
        # Add language context
        full_prompt = f"Write {language} code to: {prompt}"
        
        click.echo("\n🤖 Generating code...\n")
        generator = CodeGenerator()
        generated_code = generator.generate(full_prompt, max_length=length, temperature=temperature)
        
        click.echo("=" * 60)
        click.echo("Generated Code:")
        click.echo("=" * 60)
        click.echo(generated_code)
        click.echo("=" * 60 + "\n")
        
        # Save to file if requested
        if output:
            Path(output).write_text(generated_code)
            click.echo(f"✓ Code saved to: {output}")
    
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        exit(1)

if __name__ == '__main__':
    main()