Create app.py

app.py

@@ -0,0 +1,58 @@
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import gradio as gr
+
+MODEL_ID = "LMSeed/GPT2-small-distilled-100M"
+# HF_TOKEN = os.environ.get("HF_TOKEN")  # 如果模型私有
+
+device = 0 if torch.cuda.is_available() else -1
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
+if torch.cuda.is_available():
+    model = model.to("cuda")
+
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1
+)
+
+def chat_with_model(user_message, chat_history, max_new_tokens=60, temperature=0.8, top_p=0.9):
+    history_text = ""
+    if chat_history:
+        for i, msg in enumerate(chat_history):
+            role = "User" if i % 2 == 0 else "Assistant"
+            history_text += f"{role}: {msg}\n"
+    history_text += f"User: {user_message}\nAssistant: "
+
+    outputs = generator(history_text, max_new_tokens=max_new_tokens,
+                        do_sample=True, temperature=float(temperature),
+                        top_p=float(top_p), num_return_sequences=1)
+    reply = outputs[0]["generated_text"][len(history_text):].strip()
+    if "\n" in reply:
+        reply = reply.split("\n")[0].strip()
+    chat_history = chat_history or []
+    chat_history.append(user_message)
+    chat_history.append(reply)
+    return "", chat_history
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Chat with GPT-2")
+    with gr.Row():
+        with gr.Column(scale=3):
+            chat = gr.Chatbot(elem_id="chatbot", label="Conversation")
+            msg = gr.Textbox(label="Your message")
+            send = gr.Button("Send")
+            max_tokens = gr.Slider(10, 256, value=60, label="max_new_tokens")
+            temp = gr.Slider(0.1, 1.2, value=0.8, label="temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, label="top_p")
+        with gr.Column(scale=1):
+            gr.Markdown("Model: " + MODEL_ID)
+
+    state = gr.State([])
+    send.click(fn=chat_with_model, inputs=[msg, state, max_tokens, temp, top_p],
+               outputs=[msg, state])
+    demo.launch()