yasserrmd commited on
Commit
f2fb467
·
verified ·
1 Parent(s): d538e60

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -4
app.py CHANGED
@@ -1,10 +1,21 @@
1
  import gradio as gr
2
- from transformers import pipeline
3
 
 
 
 
 
 
 
 
 
 
 
4
  generator = pipeline(
5
- "text-generation",
6
- model="yasserrmd/SmolLM2-156M-synthetic-dlp"
7
- )
 
8
 
9
  def chat_assistant(chat_history, user_input):
10
  """Generate a response based on user input and chat history."""
 
1
  import gradio as gr
2
+ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
3
 
4
+ # Load model with 8-bit precision
5
+ model_name = "yasserrmd/SmolLM2-156M-synthetic-dlp"
6
+ model = AutoModelForCausalLM.from_pretrained(
7
+ model_name,
8
+ load_in_8bit=True, # Enable 8-bit precision
9
+ device_map="auto" # Automatically allocate to GPU/CPU
10
+ )
11
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
12
+
13
+ # Load the pipeline
14
  generator = pipeline(
15
+ "text-generation",
16
+ model=model,
17
+ tokenizer=tokenizer
18
+ )
19
 
20
  def chat_assistant(chat_history, user_input):
21
  """Generate a response based on user input and chat history."""