#
# 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}')

model = "meta-llama/Meta-Llama-3-8B-Instruct"
model = "ibm-granite/granite-3b-code-instruct"

print(f'Loading model {model}')

pipe = pipeline("text-generation", model, 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

input_textbox = gr.Textbox(
            label="Code goes here",
            info="Initial code",
            lines=3,
            value="def fibonacci(n: int) -> List:\n",
        )
gr.Interface(
    fn=generate,
    inputs=input_textbox,
    outputs=gr.Text(),
    title=model
).launch()