import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import gradio as gr
import torch

if torch.cuda.is_available():
    tokenizer = AutoTokenizer.from_pretrained("ai-forever/mGPT-13B")
    model = AutoModelForCausalLM.from_pretrained("ai-forever/mGPT-13B", load_in_8bit=True, device_map="auto")
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)

@spaces.GPU
def predict(text):
    return pipe(text)

demo = gr.Interface(
    fn=greet,
    inputs=["text"],
    outputs=["text"],
)

demo.launch()