File size: 4,079 Bytes
a675195
 
 
 
 
 
 
 
 
 
 
 
 
cf7e4de
4b7afea
cf7e4de
4b7afea
 
 
 
 
 
cf7e4de
4b7afea
 
 
 
 
 
cf7e4de
a675195
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf7e4de
 
 
 
a675195
f7d6367
 
 
a675195
 
 
 
 
 
 
f7d6367
 
a675195
 
 
 
cf7e4de
 
 
 
 
a675195
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
"""
MM Coder Agent v1 - Gradio Space
A coding assistant AI based on Qwen2.5-1.5B
"""
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from peft import PeftModel, PeftConfig
import os

# Configuration
MODEL_ID = "amkyawdev/mm-coder-agent-v1-combined"

# System prompt for better responses
SYSTEM_PROMPT = """You are a Python coding assistant. 

Output requirements:
1. Output ONLY the code, nothing else
2. Use proper Python syntax
3. Include necessary imports
4. Start with code directly in a python code block
5. Do not add explanations unless asked

Examples:
Input: print hello world
Output:
```python
print("Hello, World!")
```"""

# Global variables to store model and tokenizer
model = None
tokenizer = None

def load_model():
    """Load the model and tokenizer"""
    global model, tokenizer
    
    print("Loading model...")
    peft_config = PeftConfig.from_pretrained(MODEL_ID)
    
    base_model = AutoModelForCausalLM.from_pretrained(
        peft_config.base_model_name_or_path,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    ).eval()
    
    tokenizer = AutoTokenizer.from_pretrained(
        peft_config.base_model_name_or_path,
        trust_remote_code=True
    )
    
    model = PeftModel.from_pretrained(base_model, MODEL_ID)
    print("Model loaded successfully!")

def generate_code(prompt, max_tokens=512, temperature=0.7):
    """Generate code based on the prompt"""
    if model is None:
        load_model()
    
    # Format prompt with system instruction
    formatted_prompt = f"{SYSTEM_PROMPT}\n\nUser: {prompt}\nAssistant:"
    
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
    
    # Get pad token id safely
    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=int(max_tokens), 
            temperature=temperature,
            do_sample=True,
            top_p=0.95,
            pad_token_id=pad_token_id,
            eos_token_id=pad_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Remove the prompt and system prompt from response
    prefix_to_remove = f"{SYSTEM_PROMPT}\n\nUser: {prompt}\nAssistant:"
    if response.startswith(prefix_to_remove):
        response = response[len(prefix_to_remove):].strip()
    elif response.startswith(prompt):
        response = response[len(prompt):].strip()
    
    return response

# Create Gradio Interface
with gr.Blocks(title="🤖 MM Coder Agent v1") as demo:
    gr.Markdown("""
    # 🤖 MM Coder Agent v1
    
    A professional AI coding assistant fine-tuned from Qwen2.5-1.5B-Instruct.
    
    ## Features
    - Code Generation in multiple languages
    - Bug Detection and Fixing
    - Algorithm Implementation
    - Code Review Assistance
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Your Prompt", 
                placeholder="Write a Python function to calculate fibonacci numbers...",
                lines=5
            )
            with gr.Row():
                max_tokens = gr.Slider(64, 1024, value=512, step=64, label="Max Tokens")
                temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
            
            generate_btn = gr.Button("Generate Code", variant="primary")
        
        with gr.Column():
            output = gr.Code(label="Generated Code", language="python")
    
    generate_btn.click(
        generate_code, 
        [prompt, max_tokens, temperature], 
        output
    )
    
    gr.Markdown("""
    ---
    ### Examples
    - "Write a Python function to reverse a string"
    - "Create a React component for a login form"
    - "How to fix CORS error in Express.js?"
    - "Explain bubble sort algorithm"
    """)

# Launch the demo
demo.launch(server_name="0.0.0.0", server_port=7860)