| import gradio as gr |
| import torch |
| import os |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
| |
| LOCAL_MODEL_PATH = os.path.join(os.getcwd(), "model") |
| print("Loading model from:", LOCAL_MODEL_PATH) |
|
|
| |
| tokenizer = AutoTokenizer.from_pretrained( |
| LOCAL_MODEL_PATH, |
| local_files_only=True, |
| trust_remote_code=False |
| ) |
|
|
| model = AutoModelForCausalLM.from_pretrained( |
| LOCAL_MODEL_PATH, |
| local_files_only=True, |
| trust_remote_code=False, |
| torch_dtype=torch.float16 |
| ).to("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| |
| def chat(user_input): |
| inputs = tokenizer(user_input, return_tensors="pt").to(model.device) |
| outputs = model.generate( |
| input_ids=inputs["input_ids"], |
| attention_mask=inputs.get("attention_mask"), |
| max_new_tokens=200, |
| temperature=0.7, |
| top_p=0.9, |
| do_sample=True |
| ) |
| return tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True) |
|
|
| |
| interface = gr.Interface( |
| fn=chat, |
| inputs=gr.Textbox(lines=3, label="输入你的问题"), |
| outputs=gr.Textbox(lines=10, label="模型回答"), |
| title="我的AI模型", |
| description="在 HuggingFace Space 上运行的 LLaMA/GPT 模型网页" |
| ) |
|
|
| interface.launch() |
|
|
