DexterSptizu commited on
Commit
feb9574
1 Parent(s): b0cedb6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -0
app.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import gradio as gr
4
+
5
+ # Load the tokenizer and model
6
+ model_path = 'nvidia/Minitron-4B-Base'
7
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
8
+ model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cpu', torch_dtype=torch.float32)
9
+
10
+ def generate_text(prompt, max_length=100):
11
+ # Encode the input text
12
+ inputs = tokenizer.encode(prompt, return_tensors='pt')
13
+
14
+ # Generate the output
15
+ outputs = model.generate(
16
+ inputs,
17
+ max_length=max_length,
18
+ num_return_sequences=1,
19
+ no_repeat_ngram_size=2,
20
+ do_sample=True,
21
+ temperature=0.7,
22
+ pad_token_id=tokenizer.eos_token_id
23
+ )
24
+
25
+ # Decode and return the output
26
+ generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
27
+ return generated_text
28
+
29
+ # Create Gradio interface
30
+ demo = gr.Interface(
31
+ fn=generate_text,
32
+ inputs=[
33
+ gr.Textbox(label="Enter your prompt", placeholder="Type your prompt here..."),
34
+ gr.Slider(minimum=20, maximum=200, value=100, step=10, label="Max Length")
35
+ ],
36
+ outputs=gr.Textbox(label="Generated Text"),
37
+ title="Text Generation with Minitron-4B",
38
+ description="Enter a prompt and get AI-generated text completion.",
39
+ examples=[
40
+ ["Complete the paragraph: our solar system is"],
41
+ ["Write a short story about"],
42
+ ["Explain the concept of"]
43
+ ]
44
+ )
45
+
46
+ # Launch the application
47
+ if __name__ == "__main__":
48
+ demo.launch(share=False)