#
# 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")

print(f'HUGGINGFACE_API_TOKEN: {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']
    print(f'Response received!')
    r = response[0]['generated_text']
    return r

gr.Interface(
    fn=generate,
    inputs=gr.Text(),
    outputs=gr.Text(),
).launch()