import os
import sys
import requests
import gradio as gr 
import json
import numpy as np
from io import BytesIO
from PIL import Image

def translate(input_txt,prefix,suffix,num_beams,topk,beta):
    dictToSend = {'input':input_txt,'prefix':prefix,'suffix':suffix,'num_beams':num_beams,'alpha': (1-beta),'beta':beta,'topk':topk}
    response = requests.post('https://d7bf-141-3-25-29.ngrok-free.app/generate', json=dictToSend)
    response_data = json.loads(response.text)
    processed_text = response_data.get('text')
    processed_image_list = response_data.get('beam_img')
    processed_image_array = np.array(processed_image_list, dtype=np.uint8)
    processed_image = Image.fromarray(processed_image_array)
    return processed_text, processed_image

prefix= "<s>[INST] <<SYS>>\nYou are a professional translator from English to German. You translate the sentences to only Formal German even if the English sentence is Informal.\n<</SYS>>\nEnglish: "
#self.prefix = "Translate from English to German:\nEnglish: "
suffix = "\nGerman: "

example_formal = ["Where are you going today?",prefix,suffix,5,5,1]

prefix= "<s>[INST] <<SYS>>\nYou are a professional translator from English to German. You translate the sentences to only Informal German even if the English sentence is Formal.\n<</SYS>>\nEnglish: "
#self.prefix = "Translate from English to German:\nEnglish: "
suffix = "\nGerman: "
example_informal = ["Where are you going today?",prefix,suffix,5,5,1]

iface = gr.Interface(
        fn=translate,
        inputs=[gr.Textbox(lines=1, placeholder="Enter your English Sentences that you want to translate", label="English Sentence"), gr.Textbox(lines=5, placeholder=prefix,label="Prefix for LLM Prompt - Replace formal and informal here to control Formality Level"),gr.Textbox(lines=2, placeholder=suffix, label="DONT CHANGE!!"),gr.Slider(label='Select beam size',value=5,minimum=1,maximum=5,step=1),gr.Slider(label="TopK for rescoring",value=5,minimum=5,maximum=10,step=1),gr.Slider(label="LLM weight when rescoring, 1 for complete LLM scoring, 0 for NLLB Decoding without LLM",value=1,minimum=0,maximum=1,step=0.1)],
        outputs=[gr.Textbox(lines=1,placeholder="Enter your inputs and click submit!",label="LLM Guided NLLB German Translation"),gr.Image(label="Beam Visualization")],
        examples=[example_formal,example_informal],
        title="Ensembling NLLB 3.3B with LLama2 13B by Guided Decoding",
        )
    
iface.launch(share=True)