Detsutut commited on
Commit
ab02808
1 Parent(s): ed003fe

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import transformers
3
+ import torch
4
+
5
+ # Initialize the model
6
+ model_id = "bmi-labmedinfo/Igea-350M-v0.0.1"
7
+ pipeline = transformers.pipeline(
8
+ "text-generation",
9
+ model=model_id,
10
+ model_kwargs={"torch_dtype": torch.bfloat16},
11
+ device_map="auto",
12
+ )
13
+
14
+ # Define the function to generate text
15
+ def generate_text(input_text, max_new_tokens=128, temperature=1.0, top_k=50, top_p=0.95):
16
+ output = pipeline(
17
+ input_text,
18
+ max_new_tokens=max_new_tokens,
19
+ temperature=temperature,
20
+ top_k=top_k,
21
+ top_p=top_p,
22
+ )
23
+ return output[0]['generated_text']
24
+
25
+ # Create the Gradio interface
26
+ iface = gr.Interface(
27
+ fn=generate_text,
28
+ inputs=[
29
+ gr.inputs.Textbox(lines=2, placeholder="Enter your text here...", label="Input Text"),
30
+ gr.inputs.Slider(minimum=1, maximum=200, default=128, step=1, label="Max New Tokens"),
31
+ gr.inputs.Slider(minimum=0.1, maximum=2.0, default=1.0, step=0.1, label="Temperature"),
32
+ gr.inputs.Slider(minimum=1, maximum=100, default=50, step=1, label="Top-k"),
33
+ gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.95, step=0.01, label="Top-p")
34
+ ],
35
+ outputs="text",
36
+ title="Text Generation Interface",
37
+ description="Enter a prompt to generate text using the Igea-350M model and adjust the hyperparameters."
38
+ )
39
+
40
+ # Launch the interface
41
+ if __name__ == "__main__":
42
+ iface.launch()