rishiraj commited on
Commit
117e32e
·
verified ·
1 Parent(s): c065228

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import spaces
4
+ import torch
5
+
6
+ model_name = "sarvamai/sarvam-translate"
7
+
8
+ # Load tokenizer and model
9
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
10
+ model = AutoModelForCausalLM.from_pretrained(model_name).to('cuda:0')
11
+
12
+ @spaces.GPU
13
+ def generate(tgt_lang, input_txt):
14
+ messages = [
15
+ {"role": "system", "content": f"Translate the following sentence into {tgt_lang}."},
16
+ {"role": "user", "content": input_txt},
17
+ ]
18
+
19
+ # Apply chat template to structure the conversation
20
+ text = tokenizer.apply_chat_template(
21
+ messages,
22
+ tokenize=False,
23
+ add_generation_prompt=True
24
+ )
25
+
26
+ # Tokenize and move input to model device
27
+ model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
28
+
29
+ # Generate the output
30
+ generated_ids = model.generate(
31
+ **model_inputs,
32
+ max_new_tokens=1024,
33
+ do_sample=True,
34
+ temperature=0.01,
35
+ num_return_sequences=1
36
+ )
37
+ output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
38
+ return tokenizer.decode(output_ids, skip_special_tokens=True)
39
+
40
+ demo = gr.Interface(
41
+ fn=generate,
42
+ inputs=[
43
+ gr.Radio(["Hindi", "Bengali", "Marathi", "Telugu", "Tamil", "Gujarati", "Urdu", "Kannada", "Odia", "Malayalam", "Punjabi", "Assamese", "Maithili", "Santali", "Kashmiri", "Nepali", "Sindhi", "Dogri", "Konkani", "Manipuri (Meitei)", "Bodo", "Sanskrit"], label="Target Language", value="Hindi"),
44
+ gr.Textbox(label="Input Text", value="Be the change you wish to see in the world."),
45
+ ],
46
+ outputs=gr.Textbox(label="Translation"),
47
+ title="translate"
48
+ )
49
+ demo.launch()