playground / app.py
appvoid's picture
Update app.py
4c5c253 verified
raw
history blame
534 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-2406")
@spaces.GPU
def predict(prompt):
completion = model(prompt, max_new_tokens=64, temperature=0.2)[0]["generated_text"]
return completion
gr.Interface(
fn=predict,
inputs="text",
outputs="text",
title="playground",
).launch()