from gradio_client import Client

from deep_translator import GoogleTranslator

def predict(prompt, client, model = "0.1"):
	global iddb
	if model == "0.1":
		result = client.predict(
    	prompt,
    	"Default",
    	api_name="/chat"
		)
	elif model == "0.2":
		result = client.predict(
			prompt,
			0.3,  # 'Temperature'
			128,  # 'Max new tokens'
			0.8,  # 'Top-p (nucleus sampling)'
			1.5,  # 'Repetition penalty'
			api_name="/chat"
	  )

	return result

#  text IN  language IN
def translate(text, source):
	if source == "ru":
		target = "en"
	elif source == "en":
		target = "ru"
	out = GoogleTranslator(source = source, target = target).translate(text)
	return out


iddb = {}

def gen(text, id, model):
	global iddb

	if str(id) not in iddb:
		if model == "0.1":
			client = Client("https://afischer1985-ai-interface.hf.space/--replicas/salfk/")
		elif model == "0.2":
			client = Client("https://skier8402-mistral-super-fast.hf.space/")
		iddb[str(id)] = client
	else:
		client = iddb[str(id)]

	prompt = translate(text, "ru")
	predicted = predict(prompt, client, model).replace("</s>", "")

	return translate(predicted, "en")