import gradio as gr
# Load model directly
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Dorjzodovsuren/mongolian-gpt2")
model = AutoModelForCausalLM.from_pretrained("Dorjzodovsuren/mongolian-gpt2", from_flax=True)

generation_params = {
    "do_sample": True,
    "temperature": 0.3,
    "top_p": 0.95,
    "top_k": 40,
    "max_new_tokens": 64,
    "repetition_penalty": 2.1
}

# Create a text generation pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, **generation_params)

def text_generator(text):
    return pipe(text)[0]["generated_text"]

demo = gr.Interface(fn=text_generator, inputs="textbox", outputs="textbox")
    
if __name__ == "__main__":
    demo.launch()