File size: 1,417 Bytes
51857fd a6d3fd5 51857fd 1878582 a6d3fd5 1878582 a6d3fd5 1878582 51857fd |
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 43 44 45 46 47 48 |
import gradio
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_NAME = "arnir0/Tiny-LLM"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
def generate_text(prompt, model, tokenizer, max_length=4096, temperature=0.8, top_k=50, top_p=0.95):
inputs = tokenizer.encode(prompt, return_tensors="pt")
outputs = model.generate(
inputs,
max_length=max_length,
temperature=temperature,
top_k=top_k,
top_p=top_p,
do_sample=True
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_text
def my_inference_function(text):
prompt = f"Summary the context below\n\n{text}"
generated_text = generate_text(prompt, model, tokenizer)
return generated_text[len(prompt):]
gradio_interface = gradio.Interface(
fn=my_inference_function,
inputs="text",
outputs="text",
examples=[
["Jill"],
["Sam"]
],
title="REST API with Gradio and Huggingface Spaces",
description="This is a demo of how to build an AI powered REST API with Gradio and Huggingface Spaces – for free! Based on [this article](https://www.tomsoderlund.com/ai/building-ai-powered-rest-api). See the **Use via API** link at the bottom of this page.",
article="© Tom Söderlund 2022"
)
gradio_interface.launch()
|