Spaces:
Sleeping
Sleeping
File size: 2,417 Bytes
07e9fc6 da08414 47c889a da08414 51194d5 8f8d839 07e9fc6 51194d5 896a6d3 51194d5 896a6d3 51194d5 896a6d3 51194d5 896a6d3 f9346d5 896a6d3 51194d5 f1c133b 51194d5 f1c133b 51194d5 896a6d3 51194d5 |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
# Load tokenizer & model
#tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map="auto", # Automatically use GPU if available
torch_dtype="auto"
)
# Ensure pad token is set for safe generation
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
if getattr(model.config, "pad_token_id", None) is None:
model.config.pad_token_id = tokenizer.pad_token_id
# Create generation pipeline
#story_generator = pipeline(
#"text-generation",
#model=model,
#tokenizer=tokenizer )
# Build a text generation pipeline
generator = pipeline(
task="text-generation",
model=model,
tokenizer=tokenizer,
return_full_text=False
)
# Function to generate stories
def generate_story(prompt, max_tokens=300, temperature=0.8):
try:
chat_prompt = chat_prompt(prompt)
outputs = generator(
chat_prompt,
max_new_tokens=int(max_tokens),
temperature=float(temperature),
do_sample=True,
top_p=0.95,
top_k=50,
repetition_penalty=1.05,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id
)
return outputs[0]["generated_text"]
except Exception as e:
return f"Error during generation: {type(e).__name__}: {e}"
# Gradio UI
with gr.Blocks() as demo:
gr.Markdown("# 📖 Interactive Story Generator (TinyLlama/TinyLlama-1.1B-Chat-v1.0)")
gr.Markdown("Type a prompt and let the AI continue your story with a powerful 1.1B model.")
prompt = gr.Textbox(
label="My Story Prompt",
placeholder="e.g., In the far future, humanity discovered a hidden planet...",
lines=3
)
max_length = gr.Slider(50, 1000, value=300, step=50, label="Story Length in new tokens")
temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Creativity")
generate_btn = gr.Button("✨ Generate Story")
output = gr.Textbox(label="Generated Story", lines=20)
generate_btn.click(
fn=generate_story,
inputs=[prompt, max_length, temperature],
outputs=output
)
demo.launch() |