SimpleGPT2 / app.py
amanm10000's picture
Update app.py
44fddb6 verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
# Load the tokenizer and model
model_name = "unsloth/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Define the generation function
def generate_response(prompt):
inputs = tokenizer.encode(prompt, return_tensors="pt")
outputs = model.generate(
inputs,
max_length=512,
num_return_sequences=1,
do_sample=True,
temperature=0.7,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# Create the Gradio interface
interface = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(lines=5, placeholder="Enter your prompt here..."),
outputs=gr.Textbox(label="Generated Response"),
title="Llama-3.2-1B-Instruct Model",
description="A simple interface to interact with the Llama-3.2-1B-Instruct model.",
)
# Launch the app
if __name__ == "__main__":
interface.launch()