k3ybladewielder commited on
Commit
d0c42cc
1 Parent(s): 5a96d87

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -2,10 +2,12 @@ import gradio as gr
2
  from transformers import pipeline
3
 
4
  task = "text-generation" # Specify the task as text generation
5
- # model = "tiiuae/falcon-7b-instruct" # Use a specific model (option 1) or
6
- model = "tiiuae/falcon-40b-instruct" # Use another specific model (option 2)
7
  get_completion = pipeline(task=task, model=model, trust_remote_code=True) # Create a text generation pipeline
8
 
 
 
9
  # Define a function to format the chat history and create a prompt for the chatbot
10
  def format_chat_prompt(message, chat_history):
11
  prompt = ""
 
2
  from transformers import pipeline
3
 
4
  task = "text-generation" # Specify the task as text generation
5
+ model = "tiiuae/falcon-7b-instruct" # Use a specific model (option 1) or
6
+ # model = "tiiuae/falcon-40b-instruct" # Use another specific model (option 2)
7
  get_completion = pipeline(task=task, model=model, trust_remote_code=True) # Create a text generation pipeline
8
 
9
+ device = "cuda" if torch.cuda.is_available() else "cpu"
10
+
11
  # Define a function to format the chat history and create a prompt for the chatbot
12
  def format_chat_prompt(message, chat_history):
13
  prompt = ""