File size: 1,234 Bytes
a861494
 
 
2c8a4e3
1074868
 
2c8a4e3
be7f41a
2c8a4e3
 
a861494
 
 
 
 
a86fbe8
 
a861494
 
 
 
 
ba441e0
50bc7d1
2c3469f
12015a7
 
b243029
 
a861494
 
 
 
 
 
 
 
 
 
2c3469f
a861494
 
 
 
 
 
 
 
 
 
 
 
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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import spaces

import torch

print('torch version:', torch.__version__)

# import torch._dynamo
# torch._dynamo.config.suppress_errors = True
# torch._dynamo.disable()
# torch._dynamo.disallow_in_graph()

import gradio as gr

from transformers import AutoTokenizer, AutoModelForCausalLM

torch.set_float32_matmul_precision('high')

max_seq_length = 2048

tokenizer = AutoTokenizer.from_pretrained("ua-l/gemma-2-9b-legal-steps200-merged-16bit-uk")
model = AutoModelForCausalLM.from_pretrained(
    "ua-l/gemma-2-9b-legal-steps200-merged-16bit-uk",
    torch_dtype=torch.float16,
).to('cuda')
# compiled_model = torch.compile(model, mode="default")


print('Model dtype:', model.dtype)


@spaces.GPU
def predict(question):
    inputs = tokenizer(
    [f'''### Question:
    {question}
    
    ### Answer:
'''], return_tensors = "pt").to("cuda")

    outputs = model.generate(**inputs, max_new_tokens = 128)
    
    results = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    return results[0]

inputs = gr.Textbox(lines=2, label="Enter a question", value="Як отримати виплати ВПО?")

outputs = gr.Textbox(label="Answer")

demo = gr.Interface(fn=predict, inputs=inputs, outputs=outputs)
demo.launch()