File size: 2,395 Bytes
324bcf0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# -*- coding: utf-8 -*-
"""translation.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1bKzjrfpxJqSrJQUuGR27-kRCSndSClBo
"""


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

device = "cuda" if torch.cuda.is_available() else "cpu"

language_model_name = "Qwen/Qwen2-1.5B-Instruct"
language_model = AutoModelForCausalLM.from_pretrained(
    language_model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(language_model_name)

def process_input(input_text, action):
    if action == "Translate to Japanese":
        prompt = f"Please translate the following English text into Japanese: {input_text}"
        lang = "ja"
    elif action == "Translate to English":
        prompt = f"Please translate the following Japanese text into English: {input_text}"
        lang = "en"
    else:
        return "Invalid action. Please choose 'Translate to English' or 'Translate to Japanese'.", "error"

    messages = [
        {"role": "system", "content": "You are a helpful AI assistant for Language Translation."},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    generated_ids = language_model.generate(
        model_inputs.input_ids,
        max_new_tokens=512
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    output_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return output_text, lang

def handle_interaction(input_text, action):
    output_text, lang = process_input(input_text, action)
    return output_text

action_options = ["Translate to English", "Translate to Japanese"]

iface = gr.Interface(
    fn=handle_interaction,
    inputs=[
        gr.Textbox(label="input text"),
        gr.Dropdown(action_options, label="select action")
    ],
    outputs=[
        gr.Textbox(label="output text"),
    ],
    title="Translation App using AI",
    description="Translate input text  based on the selected Language.",
    theme= "gradio/soft"
)

if __name__ == "__main__":
    iface.launch(share=True)