import re
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    'parsak/codegen-350M-mono-lora-instruction',
)
tokenizer = AutoTokenizer.from_pretrained('Salesforce/codegen-350M-mono')

tokenizer.pad_token_id = 0  # different to <eos>
tokenizer.padding_side = "left"  # Allow batched inference

def extract_code(input_text):
    pattern = r"'''py\n(.*?)'''"
    match = re.search(pattern, input_text, re.DOTALL)
    
    if match:
        return match.group(1)
    else:
        return None  # Return None if no match is found

def generate_code(input_text):
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
    generated_ids = model.generate(input_ids, max_length=128)
    result = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return extract_code(result)

def respond(message, chat_history, additional_inputs):
    return  f"Here's an example code:\n\n```python\n{generate_code(message)}\n```" 






gr.ChatInterface(respond,
        retry_btn= gr.Button(value="Retry"), 
        undo_btn=None, clear_btn=gr.Button(value="Clear"),
        additional_inputs=[
            gr.Dropdown(["annen", "baban"])
        ]
        ).launch()