import json
import torch
from transformers import pipeline
import streamlit as st

# Load the text-generation pipeline
pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")

delimiter = "####"
system_message = f"""
You will be provided with user data. \
The user data will be delimited with \
{delimiter} characters.

Extract key information as shown in the examples shown below, to add an item in the ERPnext application. Give the output as a python dictionary object. Do not use information outside from what
is given inside the text to fill the values. Do not provide any explaination.

Example1:  
prompt: "I had made an order for 6 Units of item Logitech G15. The platform said that there were 10 Units available in stock before i made my purchase, but now it shows that the item is out of stock even though I have made a payment of 6000 towards my order.",
"item_code": "None",
"item_name": "Logitech G15",
"item_group": "None",
"stock_uom": "Unit",
"description": "None",
"standard_rate": 1000

Example2:
prompt": "Hello, I have not received my order of 5Litres Saffola Gold oil. I made an order on 07-09-2023",
"item_code": "None",
"item_name": "Saffola Gold Oil",
"item_group": "None",
"stock_uom": "Litres",
"description": "None",
"standard_rate": "None"

Example3: 
prompt": "Please add an entry of a new item in our inventory, Code IB707, and name i-ball Keyboard K-5. Current stock includes 5000 Nos of the item, grouped under Products. Price per unit is Rs. 1500.",
"item_code": "IB707",
"item_name": "i-ball Keyboard K-5",
"item_group": "Product",
"stock_uom": "Nos",
"description": "None",
"standard_rate": 1500

"""

# Create a Streamlit app
def main():
    st.title("Prompt to JSON Tool")

    st.write("This tool generates JSON output based on the user's prompt.")

    user_input = st.text_area("Enter your prompt here")

    if st.button("Generate JSON"):
        messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": f"{delimiter}{user_input}{delimiter}"},
        ]

        prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
        output_text = outputs[0]["generated_text"]

        # Parse the relevant JSON information from the output text
        json_output = output_text.split("Example1:")[-1].strip()

        # Display the JSON output on the Streamlit app
        st.write("JSON Output:")
        st.write(json_output)

        # Save the JSON output to a file
        with open("output.json", "w") as f:
            json.dump(json_output, f, indent=4)
            st.write("JSON output saved to 'output.json'.")

if __name__ == "__main__":
    main()