Spaces:
Running
on
Zero
Running
on
Zero
File size: 596 Bytes
e9a1511 7529aa7 eaa6aa4 6d028a6 cc5d713 e8b4f94 cc5d713 6d028a6 9e0fa58 cc5d713 00a0539 7529aa7 eaa6aa4 9030a60 aa41904 eaa6aa4 cc5d713 e9a1511 |
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 |
#
# Simple example.
#
import spaces
from diffusers import DiffusionPipeline
import os
import torch
from transformers import pipeline
import gradio as gr
token = os.getenv("HUGGINGFACE_API_TOKEN")
pipe = pipeline("text-generation", "meta-llama/Meta-Llama-3-8B-Instruct", torch_dtype=torch.bfloat16, device_map="auto", token=token)
# pipe.to('cuda')
@spaces.GPU
def generate(prompt):
response = pipe(prompt, max_new_tokens=512)
# r = response[0]['generated_text'][-1]['content']
return response
gr.Interface(
fn=generate,
inputs=gr.Text(),
outputs=gr.Text(),
).launch()
|