import torch
from utils_MMD import MMD_batch2
from api_init import api_init
import gradio as gr

base_model, base_tokenizer, net, feature_ref, sigma, sigma0_u, ep, loaded_model, DEVICE = api_init()

def single_instance_detection(sentence=None):

	with torch.no_grad():
		if sentence is None: 
			sentence = "Now that you've built a demo, you'll probably want to share it with others. Gradio demos can be shared in two ways: using a temporary share link or permanent hosting on Spaces."
		sentence = sentence.strip()
		inputs = base_tokenizer([sentence], padding='max_length', truncation=True, max_length=100,return_tensors="pt").to(DEVICE)
		outputs = base_model(**inputs)
		hidden_states_all = outputs[2]
		hidden_states = hidden_states_all[-1]
		token_mask_10 = inputs['attention_mask'].unsqueeze(-1)
		hidden_states_mask_10 = hidden_states * token_mask_10

		feature_for_one_sample = net(hidden_states_mask_10)
		mmd_feature_for_one_sample = MMD_batch2(torch.cat([feature_ref,feature_for_one_sample],dim=0), feature_ref.shape[0], 0, sigma, sigma0_u, ep,is_smooth=False).to('cpu')

		y_pred_loaded = loaded_model.predict(mmd_feature_for_one_sample.detach().numpy().reshape(-1, 1))
		print("y_pred_loaded:", y_pred_loaded)

		prediction = int(y_pred_loaded[0])
		if prediction == 0:
			return "Human"
		elif prediction == 1:
			return "AI"


description = "Please input a sentence in English to classify whether it's generated by a human or AI.\n\nFor example: the texts below the box are generated by ChatGPT with a prompt: Write an essay in English titled by 'A Walk in the Woods'. You can directly click on the box or copy some texts into the input field for testing."

AI_examples = [
    ["The gentle rustling of leaves, the symphony of birdsong, and the dappled sunlight filtering through the towering trees create a serene tapestry of nature's beauty. A walk in the woods is a timeless and profound experience, one that invites us to disconnect from the hustle and bustle of modern life and reconnect with the natural world. It is a journey that nourishes the soul, rejuvenates the spirit, and reminds us of our deep-rooted connection to the earth."],
    ["The woods have always held a special place in human history and imagination. They are places of mystery and enchantment, where ancient trees bear witness to centuries of change and growth. The woods are where stories of knights and dragons, fairies and gnomes, have been woven into the fabric of our culture. They are places where poets have found inspiration, artists have discovered their muse, and philosophers have contemplated the mysteries of existence."],
    ["A walk in the woods offers a unique opportunity to escape the relentless pace of modern life. As we step onto the forest path, the cares and worries of the world seem to melt away. The woods have a way of grounding us, reminding us of our place in the grand scheme of things. In the presence of towering trees that have weathered countless storms and witnessed the passage of time, our own problems and concerns can feel insignificant, allowing us to gain a broader perspective on life."],
    ["The sensory delights of the woods are unparalleled. The air is rich with the earthy scent of leaves and moss, and the soft, mossy ground beneath our feet cradles each step. The chorus of birdsong provides a soothing soundtrack, while the play of light and shadow through the foliage creates a mesmerizing dance of colors and patterns."],
    ["But perhaps the most profound aspect of a walk in the woods is the opportunity for introspection and self-discovery. In the stillness of the forest, away from the distractions of screens and schedules, we have a chance to listen to the whispers of our own thoughts. It is a time for reflection, for sorting through the jumble of ideas and emotions that often crowd our minds."]
]

# Create the Gradio interface
iface = gr.Interface(
	fn=single_instance_detection,
	inputs="text",
	outputs="text",
	title="GPT-Inspector: Text generated by Human or AI?",
	description=description,
	examples=AI_examples,
)

# Launch the Gradio interface
iface.launch()