import gradio as gr
import time
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
prompt1=''' is there any code Snippets present in the below text? Reply with YES or NO ONLY. This is very important for my career:

Text: " {} "
'''
prompt2='''Extract Code Snippets ONLY.This is very important for my career:

Text: " {} "
'''
replacements = {
    '{': 'OpenCurlyBrace',
    '}': 'CloseCurlyBrace',
    '<': 'LessThanSign'
}

model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")

def chatbot_model(m,history):
    for char, replacement in replacements.items():
        m = m.replace(char, replacement)
    # Encode the prompt and generate response
    input_ids = tokenizer.encode(prompt1.format(m) + tokenizer.eos_token, return_tensors='pt')
    outputs = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    
    # Decode the output
    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if (decoded_output=='NO'):
        yield 'NO'
    else:
        input_ids = tokenizer.encode(prompt2.format(m) + tokenizer.eos_token, return_tensors='pt')
        outputs = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    
    # Decode the output
        decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
        for char, replacement in replacements.items():
            decoded_output = decoded_output.replace(replacement, char)    
        for i in range(len(decoded_output)):
            time.sleep(0.3)
            yield decoded_output[:i+1]

iface = gr.ChatInterface(chatbot_model,title="Code Detector",undo_btn=None,theme="soft").queue()
iface.launch()