playground / app.py
appvoid's picture
Update app.py
2031352 verified
raw
history blame
525 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/arco-2")
@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()