File size: 1,301 Bytes
c14a211
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "t-tech/T-pro-it-1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto"
)

def generate_response(prompt):
    messages = [
        {"role": "system", "content": "Π’Ρ‹ T-pro, Π²ΠΈΡ€Ρ‚ΡƒΠ°Π»ΡŒΠ½Ρ‹ΠΉ ассистСнт Π² Π’-Π’Π΅Ρ…Π½ΠΎΠ»ΠΎΠ³ΠΈΠΈ. Ввоя Π·Π°Π΄Π°Ρ‡Π° - Π±Ρ‹Ρ‚ΡŒ ΠΏΠΎΠ»Π΅Π·Π½Ρ‹ΠΌ Π΄ΠΈΠ°Π»ΠΎΠ³ΠΎΠ²Ρ‹ΠΌ ассистСнтом."},
        {"role": "user", "content": prompt}
    ]
    
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=256
    )
    
    generated_ids = [
        output_ids[len(input_ids):] 
        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response

interface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="T-pro API"
)

interface.launch()