#!pip install -qqq datasets==3.5.0

import gradio as gr
import torch
from transformers import pipeline
from datasets import load_dataset

device = "cuda" if torch.cuda.is_available() else "cpu"
device

#from google.colab import userdata
#userdata.get('llama_easyread')

# map the image (description text) -> block of text

# <block>paragraph 1</block>
# <block>paragraph 2</block>


# description text of image -> (pass to embedding model) -> get vector embedding | => Compute cosine similarity -> we get similarity score 0-1 (1 means the same 0 means not the same)
# paragraph 1 -> (pass to embedding model) -> get vector embedding               |

model_id = "meta-llama/Llama-3.2-1B-Instruct"

pipe = pipeline(
    "text-generation",
    model=model_id,
    device=device,
    torch_dtype=torch.bfloat16 if "cuda" in device else torch.float32,
)

messages = [
    {"role":"system", "content": "You're a helpful EasyRead Assistant the simplifies complex documents or content. Follow the easy read guidelines. Only provide the simiplied content, for complex terms in the simplified text, always add a footnote for definitions."}
]

def add_and_generate(history, text):
    messages.append({"role":"user","content": text})
    prompt = pipe.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    # print(prompt)
    out = pipe(prompt, max_new_tokens=150, do_sample=True, temperature=0.7, top_p=0.9)
    reply = out[0]["generated_text"][len(prompt):]
    messages.append({"role":"assistant","content":reply})
    history.append((text, reply))

    return history, ""

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    txt = gr.Textbox(placeholder="Type here...")
    txt.submit(add_and_generate, [chatbot, txt], [chatbot, txt])

demo.launch(debug=True)