File size: 865 Bytes
99dc2e8 3ca8079 d6eb62f 3ca8079 5264ade 0ad4d13 5415a35 3ca8079 f5249eb 0ad4d13 fc21262 3ca8079 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
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)
@spaces.GPU(duration=600)
def predict(prompt, temperature, max_length):
return pipe(prompt, temperature=temperature, max_length=max_length, top_p=0.95, top_k=50, do_sample=True)[0]["generated_text"]
demo = gr.Interface(
fn=predict,
title="mGPT-13B Demo",
inputs=["text", gr.Slider(minimum=0.01, maximum=1.0, value=0.7, label="temperature"), gr.Slider(minimum=1, maximum=1024, value=50, label="max_length")],
outputs=["text"],
)
demo.launch() |