playground / app.py
appvoid's picture
Update app.py
9c665c5 verified
raw
history blame
529 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, temperature=0.3)[0]["generated_text"]
return completion
gr.Interface(
fn=predict,
inputs="text",
outputs="text",
title="playground",
).launch()