test / app.py
BorderCollieWei's picture
Update app.py
3e37a08 verified
raw
history blame
1.03 kB
# -*- coding: utf-8 -*-
"""Hugging Face Llama 2 App"""
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
# 加載 LLaMA-2 模型
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf")
# 定義推理函數
def generate_text(prompt):
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
inputs.input_ids,
max_length=200,
num_beams=5,
repetition_penalty=1.2,
early_stopping=True
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# 使用 Gradio 構建界面
interface = gr.Interface(
fn=generate_text,
inputs=gr.Textbox(lines=5, placeholder="Enter your prompt here..."),
outputs="text",
title="Llama 2 Text Generator",
description="Generate text using the Llama-2-13b-chat-hf model hosted on Hugging Face Spaces."
)
# 啟動應用
if __name__ == "__main__":
interface.launch()