from langchain.prompts import PromptTemplate
from langchain.llms import CTransformers
import os
import gradio as gr

def GetLlamaResponse(topic):
    llm = CTransformers(
        model_type="google-flan",
        model="tf_model.h5",
        config={"max_new_tokens": 64, "temperature": 0.75},
    )
    template = """
    You are a helpful assistant
    """
    prompt = PromptTemplate(
        input_variables=["topic", "word_count", "temperature"],
        template=template,
    )

    response = llm(
        prompt.format(
            word_count=64,
            temperature=0.4,
            topic=topic,
        )
    )

    return response

inputs_image_url = [
    gr.Textbox(type="text", label="Topic Name"),
]

outputs_result_dict = [
    gr.Textbox(type="text", label="Result"),
]

interface_image_url = gr.Interface(
    fn=GetLlamaResponse,
    inputs=inputs_image_url,
    outputs=outputs_result_dict,
    title="Text Generation",
    cache_examples=False,
)

gr.TabbedInterface(
    [interface_image_url],
    tab_names=['Some inference']
).queue().launch()