import gradio as gr
from transformers import pipeline

generator = pipeline("text-generation", model="jinv2/finetuned-zhengding-tourism", device=-1)

def predict(input_text):
    if not input_text.endswith('\n'):
        input_text += '\n'
    output = generator(
        f"用户: {input_text}",
        max_length=200,
        num_return_sequences=1,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
        pad_token_id=generator.tokenizer.pad_token_id,
        eos_token_id=generator.tokenizer.eos_token_id
    )[0]["generated_text"]
    return output

interface = gr.Interface(fn=predict, inputs="text", outputs="text", title="正定旅游助手")
interface.launch(server_name="0.0.0.0", server_port=7860)