from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from peft import PeftModel, PeftConfig
import streamlit as st

model_name = "TheBloke/Mistral-7B-Instruct-v0.2-GPTQ"

configm = AutoConfig.from_pretrained(model_name)
configm.quantization_config["disable_exllama"] = True

model = AutoModelForCausalLM.from_pretrained(model_name,
                                             trust_remote_code=False,
                                             revision="main",
                                             config=configm
                                             )

config = PeftConfig.from_pretrained("saanvi-bot/jayson")
model = PeftModel.from_pretrained(model, "saanvi-bot/jayson", peft_config = config)

# load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

"""### Use Fine-tuned Model"""

intstructions_string = f"""convert into json format \n"""

prompt_template = lambda comment: f'''[INST] {intstructions_string} \n{comment} \n[/INST]'''

model.eval()

# Streamlit interface
st.title("Text to JSON Converter")
st.write("Enter the text you want to convert to JSON format:")

# Text input from the user
user_input = st.text_area("Input text", height=200)

# Convert input text to JSON
if st.button("Convert"):
    prompt = prompt_template(user_input)

    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=200)

    json_output = tokenizer.batch_decode(outputs)[0]
    st.json(json_output)