File size: 889 Bytes
8e7f191
430146b
8e7f191
430146b
8e7f191
1cc4136
 
8e7f191
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14629e0
 
8e7f191
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch


tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")


def launch(input, history = []):
    new_user_input_ids = tokenizer.encode(
        input + tokenizer.eos_token, return_tensors="pt"
    )

    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)

    history = model.generate(
        bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id
    ).tolist()

    response = tokenizer.decode(history[0]).split("<|endoftext|>")
    response = [
        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
    ]  

    print(response)
    return response


iface = gr.Interface(launch, inputs="text", outputs="text")
iface.launch()