import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

MODEL_NAME = "NousResearch/Genstruct-7B"

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)


@spaces.GPU
def generate_text(title, content):
    msg = [{"title": title, "content": content}]
    inputs = tokenizer.apply_chat_template(msg, return_tensors="pt").cuda()
    return tokenizer.decode(model.generate(inputs, max_new_tokens=512)[0]).split(
        tokenizer.eos_token
    )[0]


examples = [
    [
        "The History of Chess",
        "Chess is a board game of strategic skill for two players, played on a chequered board with 64 squares arranged in an 8×8 grid. The game, with its origins in India around 600 AD, has evolved over centuries to become one of the world's most popular board games. Each player begins with 16 pieces: one king, one queen, two rooks, two knights, two bishops, and eight pawns. The objective of the game is to 'checkmate' the opponent's king, whereby the king is under immediate attack (in 'check') and there is no way for it to escape.",
    ],
    [
        "Climate Change: Causes and Effects",
        "Climate change refers to long-term shifts in temperatures and weather patterns. These shifts may be natural, such as through variations in the solar cycle. But since the 1800s, human activities have been the main driver of climate change, primarily due to burning fossil fuels like coal, oil and gas. Burning fossil fuels generates greenhouse gas emissions that act like a blanket wrapped around the Earth, trapping the sun's heat and raising temperatures. Examples of greenhouse gas emissions that are causing climate change include carbon dioxide and methane.",
    ],
    [
        "The Invention of the Internet",
        "The Internet is a global system of interconnected computer networks that use the standard Internet protocol suite (TCP/IP) to link devices worldwide. It is a network of networks that consists of private, public, academic, business, and government networks of local to global scope, linked by a broad array of electronic, wireless, and optical networking technologies. The origins of the Internet date back to research commissioned by the United States government in the 1960s to build robust, fault-tolerant communication with computer networks.",
    ],
]


description = """
<h2>Welcome to the Genstruct-7B Text Generation Demo!</h2>

<p>This demo showcases the Genstruct-7B model, an innovative instruction-generation AI designed to create valid instructions from raw text. Here's what you can do:</p>

<ul>
  <li>Input a title and content in the style of a Wikipedia article or informative blog post.</li>
  <li>See the model generated questions and detailed reasoning based on your input.</li>
  <li>Explore how Genstruct-7B can be used to create partially synthetic instruction datasets from any raw text corpus.</li>
</ul>

<p>Genstruct-7B stands out by grounding its generations in user-provided context, allowing for the creation of questions involving complex scenarios that require step-by-step reasoning. Give it a try and see how it can enhance your text generation tasks!</p>
"""


demo = gr.Interface(
    fn=generate_text,
    inputs=[gr.Textbox(label="Title"), gr.Textbox(label="Content", lines=5)],
    outputs=gr.Textbox(label="Generated Output", lines=10),
    title="Genstruct-7B Text Generation Demo",
    examples=examples,
    description=description,
)

if __name__ == "__main__":
    demo.queue(max_size=10)
    demo.launch()