ameursg commited on
Commit
f2522be
Β·
verified Β·
1 Parent(s): ef4ad31

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -8
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import os
2
  import gradio as gr
3
- from transformers import pipeline
4
  from huggingface_hub import login
5
 
6
  # Get Hugging Face token securely from environment variables
@@ -12,13 +12,12 @@ if HF_TOKEN:
12
  else:
13
  raise ValueError("🚨 Hugging Face token not found! Please add it to Secrets in your Hugging Face Space.")
14
 
15
- # Load AI model (Optimized for Hugging Face GPU)
16
- chatbot = pipeline(
17
- "text-generation",
18
- model="TheBloke/Mistral-7B-Instruct-v0.1-GGUF", # Updated to GGUF version
19
- token=HF_TOKEN,
20
- device=0 # Ensures it runs on GPU
21
- )
22
 
23
  def generate_itinerary(destination, start_date, end_date, traveler_type, companion):
24
  """Generates an AI-powered travel itinerary based on user preferences"""
 
1
  import os
2
  import gradio as gr
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
4
  from huggingface_hub import login
5
 
6
  # Get Hugging Face token securely from environment variables
 
12
  else:
13
  raise ValueError("🚨 Hugging Face token not found! Please add it to Secrets in your Hugging Face Space.")
14
 
15
+ # Load AI model (Optimized for T4 GPU)
16
+ model_name = "TheBloke/Mistral-7B-Instruct-v0.1-AWQ"
17
+ tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
18
+ model = AutoModelForCausalLM.from_pretrained(model_name, token=HF_TOKEN, device_map="auto")
19
+
20
+ chatbot = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
 
21
 
22
  def generate_itinerary(destination, start_date, end_date, traveler_type, companion):
23
  """Generates an AI-powered travel itinerary based on user preferences"""