TeLLMyStory/story-generation-docker

1#last version of app.py2from transformers import AutoTokenizer, AutoModelForCausalLM, GPTQConfig3import torch4import optimum5import auto_gptq6import gradio as gr7import time8 9device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")10 11model_name = "TheBloke/zephyr-7B-beta-GPTQ"12 13tokenizer = AutoTokenizer.from_pretrained(model_name,use_fast=True,padding_side="left")14quantization_config_loading = GPTQConfig(15                                bits=4,16                                group_size=128,17                                disable_exllama=False)18model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config_loading, device_map="auto")19model = model.to(device)20 21def generate_text(input_text,max_new_tokens=512,top_k=50,top_p=0.95,temperature=0.7,no_grad=False):22    tokenizer.pad_token_id = tokenizer.eos_token_id23    input_ids = tokenizer.encode(input_text, padding=True, return_tensors="pt").to(device)24    attention_mask = input_ids.ne(tokenizer.pad_token_id).long().to(device)25    output = None26    if no_grad:27        with torch.no_grad():28            output = model.generate(input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, top_k=top_k, top_p=top_p, temperature=temperature,do_sample=True)29    else:30        output = model.generate(input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, top_k=top_k, top_p=top_p, temperature=temperature,do_sample=True)31    return tokenizer.decode(output[0], skip_special_tokens=True)32 33 34time_story = 035 36def generate_response(input,history: list[tuple[str, str]],max_tokens, temperature, top_p):37    messages=[]38    for val in history:39        # Directly access content using "content" key40        messages.extend([{"role": "user", "content": val.get("content")}, {"role": "assistant", "content": val.get("content")}]) if val else None41 42    messages.append({"role": "user", "content": input})43    44    start = time.time()45    output = generate_text(input,max_new_tokens=max_tokens, top_p=top_p, temperature=temperature)46    end = time.time()47    time_story= end-start48    print(f'Time to generate the story: {time_story}')49    history.append((input,output))50    yield output51 52#define the chatinterface53title = "TeLLMyStory"54description = "A LLM for stories generation aiming the reinforcement of the controllability aspect"55theme = gr.Theme.from_hub("Yntec/HaleyCH_Theme_Yellow_Blue")56examples=[["Once upon a time a witch named Malefique was against the wedding of her daughter with the son of the king of the nearby kingdom."],57        ["Once upon a time an ice-cream met a spoon and they fell in love"],58        ["The neverending day began with a beautiful sunshine and an AI robot which was seeking humans on the desert Earth."]]59 60demo = gr.ChatInterface(61      generate_response,62      type="messages",63      title=title,64      description=description,65      theme=theme,66      examples=examples,67      additional_inputs=[68          gr.Slider(minimum=1, maximum=2048, value=100, step=1, label="Max new tokens"),69          gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),70          gr.Slider(71              minimum=0.1,72              maximum=1.0,73              value=0.95,74              step=0.05,75              label="Top-p (nucleus sampling)",76          ),77      ],78    79    stop_btn="Stop",80    delete_cache=[60,60],81    show_progress="full",82    save_history=True,83  )84 85 86if __name__ == "__main__":87    demo.launch(share=True,debug=True)