playground / app.py
appvoid's picture
Update app.py
6147333 verified
raw
history blame
512 Bytes
import os
from threading import Thread
from typing import Iterator
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import pipeline
model = pipeline("text-generation" , model="appvoid/palmer-004")
@spaces.GPU
def predict(prompt):
completion = model(prompt, max_new_tokens=64)[0]["generated_text"]
return completion
gr.Interface(
fn=predict,
inputs="text",
outputs="text",
title="playground",
).launch()