ksh-nyp commited on
Commit
b5ae183
1 Parent(s): c5d9aa0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -1,9 +1,9 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
- # Initialize the pipeline with the model for CPU usage
5
  model_name = "ksh-nyp/llama-2-7b-chat-TCMKB2"
6
- pipe = pipeline("text-generation", model=model_name, device=0) # device=0 for CPU
7
 
8
  def generate_text(prompt):
9
  # Generate text based on the input prompt
@@ -19,5 +19,5 @@ interface = gr.Interface(
19
  description="Enter a prompt to generate text using the TCM Fine-Tuned LLaMA 2 7B model."
20
  )
21
 
22
- # Launch the app
23
  interface.launch()
 
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
+ # Initialize the pipeline with the model
5
  model_name = "ksh-nyp/llama-2-7b-chat-TCMKB2"
6
+ pipe = pipeline("text-generation", model=model_name, device=0)
7
 
8
  def generate_text(prompt):
9
  # Generate text based on the input prompt
 
19
  description="Enter a prompt to generate text using the TCM Fine-Tuned LLaMA 2 7B model."
20
  )
21
 
22
+ # Launch the app
23
  interface.launch()