import os
import gradio as gr
from langchain.llms import HuggingFaceHub

model_repo = os.getenv('HF_MODEL_REPO')
template = """[INST]<<SYS>>I want you to act as language translator. You do translate {source} texts in the context into {target} then you return to me the whole translated context AND NOTHING ELSE.<</SYS>>
Begin of the context:
{query}
End of the context.[/INST]
{target} translation of the context:

"""

model_kwargs={
            "max_new_tokens":2048,
            "temperature": 0.01,
            "truncate": 4096,
            "seed" : 256,
            "stop" : ["</s>","<|endoftext|>","<|end|>"],
            }

llm = HuggingFaceHub(repo_id=model_repo, task="text-generation", model_kwargs=model_kwargs)

def translation(source, target, text):
    input_prompt = template.replace("{source}", source)
    input_prompt = input_prompt.replace("{target}", target)
    input_prompt = input_prompt.replace("{query}", text)
    print(input_prompt)
    response=llm(input_prompt)
    return response

gr.Interface(translation, inputs=["text","text","text"], outputs="text").launch()