File size: 1,353 Bytes
686678a 951adef 26a2377 b74dff6 26a2377 951adef 686678a 26a2377 951adef 686678a 951adef 686678a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import spaces
import gradio as gr
# from airllm import HuggingFaceModelLoader, AutoModelForCausalLM
from airllm import AutoModel
import mlx.core as mx
model = AutoModel.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
# model = AutoModel.from_pretrained(model_loader)
MAX_LENGTH = 128
@spaces.GPU
def generate_text(input_text):
input_tokens = model.tokenizer(input_text,
return_tensors="np",
return_attention_mask=False,
truncation=True,
max_length=MAX_LENGTH,
padding=False)
output = model.generate(mx.array(input_tokens['input_ids']),
max_new_tokens=20,
use_cache=True,
return_dict_in_generate=True)
# input_ids = model.tokenizer.encode(input_text, return_tensors="np")
# output = model.generate(input_ids, max_length=100)
# return model.tokenizer.decode(output[0])
return output
iface = gr.Interface(
fn=generate_text,
inputs=gr.Textbox(placeholder="Enter prompt..."),
outputs="text",
title="LLaMA 3 70B Text Generation"
)
iface.launch(server_name="0.0.0.0", server_port=7860) |