rish13 commited on
Commit
2a4edc1
·
verified ·
1 Parent(s): 892e1e5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -3
app.py CHANGED
@@ -1,12 +1,16 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
- # Load the model
5
- model = pipeline("text-generation", model="rish13/polymers")
 
 
 
 
6
 
7
  def generate_response(prompt):
8
  # Generate text from the model
9
- response = model(prompt, max_length=150, num_return_sequences=1)
10
 
11
  # Get the generated text from the response
12
  generated_text = response[0]['generated_text']
 
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
+ # Load the model with optimization settings
5
+ model = pipeline(
6
+ "text-generation",
7
+ model="rish13/polymers",
8
+ device=0 # Use device=0 for GPU (if available), -1 for CPU
9
+ )
10
 
11
  def generate_response(prompt):
12
  # Generate text from the model
13
+ response = model(prompt, max_length=100, num_return_sequences=1, temperature=0.7)
14
 
15
  # Get the generated text from the response
16
  generated_text = response[0]['generated_text']