from transformers import pipeline, set_seed
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import streamlit as st

# tokenizer = AutoTokenizer.from_pretrained("microsoft/GODEL-v1_1-large-seq2seq")
# model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/GODEL-v1_1-large-seq2seq")

generator = pipeline('text-generation', model='gpt2')
set_seed(42)

def generate(instruction, knowledge, dialog):
    if knowledge != '':
        knowledge = '[KNOWLEDGE] ' + knowledge
    dialog = ' EOS '.join(dialog)
    query = f"{instruction} [CONTEXT] {dialog} {knowledge}"
    input_ids = tokenizer(f"{query}", return_tensors="pt").input_ids
    outputs = model.generate(input_ids, max_length=128, min_length=8, top_p=0.9, do_sample=True)
    output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return output

text_init = st.text_area('Enter init description')
text_deploy = st.text_area('Enter deploy description')
text_invariant = st.text_area('Enter invariant description')
if st.button('GENERATE RESPONSE'):
    # st.write(text_init, text_deploy, text_invariant)
    # Instruction
    
    # instruction = f'Using the descriptions, write the init code, deploy code and invariant code'
    # knowledge = "EXAMPLE:{init description: [“Create a new address called “user”. Give the “user” “AMOUNT” of LP tokens. Using the “user” as the sender approve the “StaxLPStaking” contract to use his LP tokens and stake all the LP tokens into the “StaxLPStaking” contract.“],\n deploy description: [“Create a new address called “owner”. Using the “owner” as the sender deploy the “StaxLP” contract and the “StaxLPStaking” contract.“],\n invariant description: [“Using the “user” address as sender withdraw all the LP tokens from the “StaxLPStaking” contract and check the balance of the LP tokens of the “user” is not less than “AMOUNT”.“],\n init code: [{\naddress public user;\nuint256 public constant AMOUNT = 10e18;\nfunction init() public {\n    user = makeAddr(“USER”);\n    deal(address(lp), user, AMOUNT);\n    vm.startPrank(user);\n    lp.approve(address(lpStaking), lp.balanceOf(user));\n    lpStaking.stakeAll();\n    vm.stopPrank();\n}\n    }],\n deploy code: [{\nStaxLP public lp;\nStaxLPStaking public lpStaking;\naddress public owner;\nfunction deploy() public {\n    owner = makeAddr(“OWNER”);\n    vm.startPrank(owner);\n    lp = new StaxLP(“Stax Frax/Temple LP Token”, “xFraxTempleLP”);\n    lpStaking = new StaxLPStaking(address(lp), owner);\n    vm.stopPrank();\n}\n    }],\n invariant code: [{\nfunction invariant() public {\n    vm.prank(user);\n    lpStaking.withdrawAll(false);\n    assert(lp.balanceOf(user) >= AMOUNT);\n}\n    }],\n}\n]"
    # dialog = [f'{{init description: [“{text_init}“],\n deploy description: [“{text_deploy}“],\n invariant description: [“{text_invariant}“],\n}}']
    # response = generate(instruction, knowledge, dialog)
    PROMPT_TEXT = f'Using the EXAMPLE1 below, write the init(), deploy() and invariant() code for the EXAMPLE2:\nEXAMPLE1:\n[\n{{\n    init description: [“Create a new address called “user”. Give the “user” “AMOUNT” of LP tokens. Using the “user” as the sender approve the “StaxLPStaking” contract to use his LP tokens and stake all the LP tokens into the “StaxLPStaking” contract.“],\n    deploy description: [“Create a new address called “owner”. Using the “owner” as the sender deploy the “StaxLP” contract and the “StaxLPStaking” contract.“],\n    invariant description: [“Using the “user” address as sender withdraw all the LP tokens from the “StaxLPStaking” contract and check the balance of the LP tokens of the “user” is not less than “AMOUNT”.“],\n    init code: [{{\naddress public user;\nuint256 public constant AMOUNT = 10e18;\nfunction init() public {{\n    user = makeAddr(“USER”);\n    deal(address(lp), user, AMOUNT);\n    vm.startPrank(user);\n    lp.approve(address(lpStaking), lp.balanceOf(user));\n    lpStaking.stakeAll();\n    vm.stopPrank();\n}}\n    }}],\n    deploy code: [{{\nStaxLP public lp;\nStaxLPStaking public lpStaking;\naddress public owner;\nfunction deploy() public {{\n    owner = makeAddr(“OWNER”);\n    vm.startPrank(owner);\n    lp = new StaxLP(“Stax Frax/Temple LP Token”, “xFraxTempleLP”);\n    lpStaking = new StaxLPStaking(address(lp), owner);\n    vm.stopPrank();\n}}\n    }}],\n    invariant code: [{{\nfunction invariant() public {{\n    vm.prank(user);\n    lpStaking.withdrawAll(false);\n    assert(lp.balanceOf(user) >= AMOUNT);\n}}\n    }}],\n}}\n]\nEXAMPLE2:\n{{\n    init description: [“{text_init}“],\n    deploy description: [“{text_deploy}“],\n    invariant description: [“{text_invariant}“],\n}}'
    responses = generator(PROMPT_TEXT, max_length=1500, num_return_sequences=5)
    for response in responses:
        st.write(response)