Spaces:
Sleeping
Sleeping
File size: 4,502 Bytes
1b7f60b |
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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import gradio as gr
import torch # Optional, but good practice if using a PyTorch model
from transformers import AutoModelForCausalLM, AutoTokenizer
# --- 1. Load a simple, small pre-trained LLM and its tokenizer ---
# We'll use DistilGPT2 for speed and small size.
# You can replace this with another small model if you prefer.
model_name = "distilgpt2"
try:
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# If you have a GPU, uncomment the next line
# model.to("cuda" if torch.cuda.is_available() else "cpu")
model_loaded = True
print(f"Successfully loaded model and tokenizer for: {model_name}")
except Exception as e:
print(f"Error loading model: {e}")
model_loaded = False
# Define dummy functions if model fails to load, so Gradio interface still launches
def generate_text_from_llm(prompt_text):
return "Error: Model could not be loaded. Please check server logs."
tokenizer = None # To avoid errors later if tokenizer specific functions are called
if model_loaded and tokenizer:
# Ensure pad_token is set if it's not already (important for generate)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id
# --- 2. Define the LLM inference function ---
def generate_text_from_llm(prompt_text):
"""
Generates a short text continuation using the loaded LLM.
"""
if not prompt_text:
return "Please enter a starting prompt!"
try:
# Encode the input prompt
inputs = tokenizer.encode(prompt_text, return_tensors="pt", truncation=True, max_length=512)
# If you have a GPU, uncomment the next line
# inputs = inputs.to("cuda" if torch.cuda.is_available() else "cpu")
# Generate text
# max_length is the total length of prompt + generated text
# num_return_sequences=1 means we want one completion
# no_repeat_ngram_size helps avoid repetitive text
outputs = model.generate(
inputs,
max_length=len(inputs[0]) + 50, # Generate up to 50 new tokens
num_return_sequences=1,
pad_token_id=tokenizer.eos_token_id, # Use EOS token for padding during generation
no_repeat_ngram_size=2, # Avoid repeating 2-grams
early_stopping=True
)
# Decode the generated text
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Return only the newly generated part (optional, can be tricky)
# For simplicity, we'll return the whole thing for now.
# To return only new text: return generated_text[len(prompt_text):].strip()
return generated_text
except Exception as e:
print(f"Error during generation: {e}")
return f"Error during text generation: {e}"
# --- 3. Create the Gradio Interface ---
demo = gr.Interface(
fn=generate_text_from_llm,
inputs=[
gr.Textbox(
label="Enter your prompt",
placeholder="Start typing here...",
lines=5
)
],
outputs=[
gr.Textbox(label="LLM Generated Text", lines=10)
],
title="📝 Simple LLM Text Generator",
description="Enter a prompt and a small LLM (DistilGPT2) will try to continue it. This is a basic demo for learning purposes.",
examples=[
["Once upon a time, in a land far away,"],
["The best way to learn programming is"],
["Artificial intelligence is rapidly changing the world by"]
],
theme=gr.themes.Soft() # You can try other themes like gr.themes.Default()
)
# --- 4. Launch the app ---
# When deploying to Hugging Face Spaces, they will run this launch() command.
# For local testing with a shareable link, use share=True.
if __name__ == "__main__":
if model_loaded:
demo.launch(debug=True, share=True) # share=True creates a temporary public link
else:
print("Model failed to load. Gradio app will run with an error message function.")
# Launch with the dummy function so the UI still appears
demo_error = gr.Interface(fn=lambda x: "Error: Model could not be loaded.", inputs="textbox", outputs="textbox", title="LLM Demo - MODEL LOAD ERROR")
demo_error.launch(debug=True, share=True) |