|
import os |
|
import spaces |
|
import torch |
|
import transformers |
|
import gradio as gr |
|
|
|
access_token = os.environ["HF_ACCESS_TOKEN"] |
|
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct" |
|
|
|
pipeline = transformers.pipeline( |
|
"text-generation", |
|
model=model_id, |
|
model_kwargs={"torch_dtype": torch.bfloat16}, |
|
device_map="auto", |
|
token=access_token |
|
) |
|
|
|
@spaces.GPU |
|
def generate_text(input_text): |
|
|
|
prompt = {"role": "system", "content": "Summarize the following: "+input_text} |
|
output = pipeline(prompt, |
|
max_new_tokens=256, |
|
) |
|
return output |
|
|
|
|
|
iface = gr.Interface( |
|
fn=generate_text, |
|
inputs=gr.Textbox(placeholder="Enter prompt..."), |
|
outputs="text", |
|
title="LLaMA 3 8B Text Generation" |
|
) |
|
|
|
iface.launch(server_name="0.0.0.0", server_port=7860) |
