# Available on https://huggingface.co/spaces/CedricZ/Llama3_8B_Instruct_Translator

import gradio as gr
import time
from utils import format_as_translator
import requests

def chatbot_demo(message, history, languages):

    #We don't concatenate history data to avoid noisy data, and avoid the input message becomes too long
    input_message = format_as_translator(message, languages)
    
    #Add another assistant delimiter at begining to make sure the output text doesn't contain 'assistant/n/n'
    json_obj = {
    "inputs": input_message + '<|eot_id|><|start_header_id|>translator<|end_header_id|>\n\n',
    "parameters": {
      "best_of": 1,
      "decoder_input_details": False,
      "details": True,
      "do_sample": True,
      "frequency_penalty": 0.1,
      "grammar": None,
      "max_new_tokens": 500,
      "repetition_penalty": 1.03,
      "return_full_text": False,
      "seed": None,
      "stop": [
        "photographer"
      ],
      "temperature": 0.5,
      "top_k": 1,
      "top_n_tokens": 5,
      "top_p": 0.95,
      "truncate": None,
      "typical_p": 0.95,
      "watermark": True
    }
  }
    
    response = requests.post('https://uf9t072wj5ki2ho4.eu-west-1.aws.endpoints.huggingface.cloud/generate', json=json_obj)
    data = response.json()
    llama_out = data['generated_text']
    for i in range(len(llama_out)):
        time.sleep(0.05)
        yield llama_out[: i + 1]

demo = gr.ChatInterface(
    fn=chatbot_demo,
    additional_inputs=[gr.Textbox("German", label="Target Language", info='SOURCE LANGUAGE SHOULD BE IN ENG!')],
    chatbot=gr.Chatbot(height=500),
    textbox=gr.Textbox(placeholder="Just ask Llama3 to translate your sentences into any language you want!", container=False, scale=15),
    cache_examples=False,
    title="Llama 3 8B Instruct - Prompted Translator",
    )

if __name__ == "__main__":
    demo.launch()