File size: 1,962 Bytes
33d6f68 b88e2df 33d6f68 |
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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer, GPT2Config
from rotary_embedding_torch import RotaryEmbedding
import gradio as gr
import spaces
# Define the max length used during training
max_length = 8192
# Load the model and tokenizer from Hugging Face Hub
model_name = "archit11/gpt2-long-finetuned"
config = GPT2Config.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# Add rotary embeddings
rotary_emb = RotaryEmbedding(
dim=32,
interpolate_factor=4.0,
)
for layer in model.transformer.h:
layer.attn.rotary_emb = rotary_emb
# Set the model to evaluation mode
model.eval()
# Define the inference function
@spaces.GPU(duration=120)
def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9):
input_ids = tokenizer.encode(prompt, return_tensors="pt")
# Generate text
with torch.no_grad():
output = model.generate(
input_ids,
max_length=max_length,
temperature=temperature,
top_p=top_p,
num_return_sequences=1,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
return generated_text
# Create Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(lines=5, label="Prompt"),
gr.Slider(minimum=10, maximum=1000, value=100, step=10, label="Max Length"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p")
],
outputs=gr.Textbox(lines=10, label="Generated Text"),
title="Custom GPT-2 Text Generation",
description="Enter a prompt to generate text using the custom-trained GPT-2 model."
)
# Launch the interface
iface.launch() |