app.py

import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer


def init_model():
    model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0", torch_dtype=torch.float16, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=False, trust_remote_code=True)
    return model, tokenizer
    


def chat(prompt, top_k, temperature):
    prompt = f"### Instruction:{prompt.strip()}  ### Response:"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
    generate_ids = model.generate(inputs.input_ids, do_sample=True, max_new_tokens=2048, top_k=int(top_k), top_p=0.84, temperature=float(temperature), repetition_penalty=1.15, eos_token_id=2, bos_token_id=1, pad_token_id=0)
    response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    response = response.lstrip(prompt)
    print('-log: ',prompt, response)
    return response


if __name__ == '__main__':
    model, tokenizer = init_model()
    demo = gr.Interface(
        fn=chat,
        inputs=["text", gr.Slider(1, 60, value=10, step=1), gr.Slider(0.1, 2.0, value=1.0, step=0.1)],
        outputs="text",
    )
    demo.launch()