File size: 2,152 Bytes
1baa2a7 d0f3977 d101153 d0f3977 d8cb255 d0f3977 d101153 d0f3977 d8cb255 03542e5 510e759 d0f3977 510e759 d101153 af494ea 510e759 d0f3977 1baa2a7 d0f3977 d8cb255 |
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 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load your hosted model and tokenizer from Hugging Face.
model_name = "Samurai719214/gptneo-mythology-storyteller"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Use GPU if available.
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
def generate_full_story(excerpt: str) -> str:
"""
Given an incomplete story excerpt (without header details), this function calls the model
to generate the complete story that includes Parv, Key Event, Section and the story continuation.
"""
# Tokenize the user-provided excerpt.
encoded_input = tokenizer(excerpt, return_tensors = "pt")
# Move tensors to the appropriate device.
encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
# Generate tokens. Here, we set parameters to control length and creativity.
output = model.generate(
encoded_input["input_ids"],
attention_mask = encoded_input["attention_mask"],
max_new_tokens = 200, # Generate 200 new tokens on top of the input.
do_sample = True,
temperature = 0.8,
top_p = 0.95,
no_repeat_ngram_size = 2,
return_dict_in_generate = True
)
# Decode the generated sequence.
generated_text = tokenizer.decode(output.sequences[0], skip_special_tokens = True)
return generated_text
# Build the Gradio interface.
interface = gr.Interface(
fn = generate_full_story,
inputs = gr.Textbox(
lines = 5,
label = "Incomplete story excerpt",
placeholder = "Enter an excerpt from the Mahabharata here..."
),
outputs = gr.Textbox(label = "Chapter summary"),
title = "🏺 Mythology Storyteller",
description = (
"Enter a phrase from a chapter of your choice (if possible please enter the Parv, Key Event, & Section for an accurate answer). "
"The model will generate the summary of the respective chapter."
)
)
# Launch the Gradio app.
interface.launch(share = True) |