app.py

import os
import gradio as gr
import torch
import json
from transformers import LlamaTokenizer, LlamaForCausalLM, LlamaConfig
from peft import PeftModel

# Set Hugging Face Token for Authentication
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # Ensure this is set in your environment

if not HUGGINGFACE_TOKEN:
    raise ValueError("❌ HUGGINGFACE_TOKEN is not set. Please set it in your environment.")

print("✅ HUGGINGFACE_TOKEN is set.")

# Model Paths
MODEL_PATH = "meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8"  # Directly using quantized model
LLAMA_GUARD_NAME = "meta-llama/Llama-Guard-3-1B-INT4"

# Function to load Llama model (without LoRA)
def load_quantized_model(model_path):
    print(f"🔄 Loading Quantized Model: {model_path}")

    # Load the config manually
    config = LlamaConfig.from_pretrained(model_path)

    # Initialize model
    model = LlamaForCausalLM(config)

    # Load the quantized weights manually
    checkpoint_path = os.path.join(model_path, "consolidated.00.pth")
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"❌ Checkpoint file not found: {checkpoint_path}")

    state_dict = torch.load(checkpoint_path, map_location="cpu")

    # Load the state dict into the model
    model.load_state_dict(state_dict, strict=False)

    # Move model to GPU if available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)

    print("✅ Quantized model loaded successfully!")
    return model

# Load Tokenizer
tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH, token=HUGGINGFACE_TOKEN, legacy=False)

# Load the model
model = load_quantized_model(MODEL_PATH)

# Load the quantized Llama model
tokenizer, model = load_llama_model(QUANTIZED_MODEL)

# Load Llama Guard for content moderation
guard_tokenizer, guard_model = load_llama_model(LLAMA_GUARD_NAME)

# Define Prompt Templates
PROMPTS = {
    "project_analysis": """Analyze this project description and generate:
1. Project timeline with milestones
2. Required technology stack
3. Potential risks
4. Team composition
5. Cost estimation
Project: {project_description}""",
    
    "code_generation": """Generate implementation code for this feature:
{feature_description}
Considerations:
- Use {programming_language}
- Follow {coding_standards}
- Include error handling
- Add documentation""",

    "risk_analysis": """Predict potential risks for this project plan:
{project_data}
Format output as JSON with risk types, probabilities, and mitigation strategies"""
}

# Function: Content Moderation using Llama Guard
def moderate_input(user_input):
    prompt = f"""Input: {user_input}
Please verify that this input doesn't violate any content policies."""

    inputs = guard_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)

    with torch.no_grad():
        outputs = guard_model.generate(inputs.input_ids, max_length=256, temperature=0.1)
    
    response = guard_tokenizer.decode(outputs[0], skip_special_tokens=True)

    if any(flag in response.lower() for flag in ["flagged", "violated", "policy violation"]):
        return "⚠️ Content flagged by Llama Guard. Please modify your input."
    
    return None

# Function: Generate AI responses
def generate_response(prompt_type, **kwargs):
    prompt = PROMPTS[prompt_type].format(**kwargs)
    
    moderation_warning = moderate_input(prompt)
    if moderation_warning:
        return moderation_warning

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)

    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_length=512,
            temperature=0.7 if prompt_type == "project_analysis" else 0.5,
            top_p=0.9,
            do_sample=True
        )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Define UI functions
def analyze_project(project_description):
    return generate_response("project_analysis", project_description=project_description)

def generate_code(feature_description, programming_language, coding_standards):
    return generate_response("code_generation", feature_description=feature_description, programming_language=programming_language, coding_standards=coding_standards)

def predict_risks(project_data):
    return generate_response("risk_analysis", project_data=project_data)

# Gradio UI Setup
def create_gradio_interface():
    with gr.Blocks(title="AI Project Manager", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🚀 AI-Powered Project Manager & Code Assistant")
        
        with gr.Tab("Project Setup"):
            project_input = gr.Textbox(label="Project Description", lines=5, placeholder="Describe your project...")
            project_output = gr.Textbox(label="Project Analysis", lines=15)
            analyze_btn = gr.Button("Analyze Project")
            analyze_btn.click(analyze_project, inputs=project_input, outputs=project_output)
        
        with gr.Tab("Code Assistant"):
            code_input = gr.Textbox(label="Feature Description", lines=3)
            lang_select = gr.Dropdown(["Python", "JavaScript", "Java", "C++"], label="Language", value="Python")
            standards_select = gr.Dropdown(["PEP8", "Google", "Airbnb"], label="Coding Standard", value="PEP8")
            code_output = gr.Code(label="Generated Code")
            code_btn = gr.Button("Generate Code")
            code_btn.click(generate_code, inputs=[code_input, lang_select, standards_select], outputs=code_output)
        
        with gr.Tab("Risk Analysis"):
            risk_input = gr.Textbox(label="Project Plan", lines=5)
            risk_output = gr.JSON(label="Risk Predictions") 
            risk_btn = gr.Button("Predict Risks")
            risk_btn.click(predict_risks, inputs=risk_input, outputs=risk_output)
        
        with gr.Tab("Live Collaboration"):
            gr.Markdown("## Real-time Project Collaboration")
            chat = gr.Chatbot(height=400)
            msg = gr.Textbox(label="Chat with AI PM")
            clear = gr.Button("Clear Chat")
            
            def respond(message, chat_history):
                moderation_warning = moderate_input(message)
                if moderation_warning:
                    chat_history.append((message, moderation_warning))
                    return "", chat_history

                history_text = ""
                for i, (usr, ai) in enumerate(chat_history[-3:]):
                    history_text += f"User: {usr}\nAI: {ai}\n"
                
                prompt = f"""Project Management Chat:
Context: {message}
Chat History: {history_text}
User: {message}"""

                inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
                
                with torch.no_grad():
                    outputs = model.generate(
                        inputs.input_ids,
                        max_length=1024,
                        temperature=0.7,
                        top_p=0.9,
                        do_sample=True
                    )
                
                response = tokenizer.decode(outputs[0], skip_special_tokens=True)
                chat_history.append((message, response))
                return "", chat_history
            
            msg.submit(respond, [msg, chat], [msg, chat])
            clear.click(lambda: None, None, chat, queue=False)

    return demo

# Run Gradio App
if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch(share=True)