Update app.py

app.py

@@ -7,6 +7,7 @@ import sqlparse
 import logging
 
 # Import necessary modules from transformers and langchain
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain.llms import HuggingFacePipeline
 
@@ -24,14 +25,20 @@ if hf_token is None:
     st.error("Hugging Face API token is not set. Please set the HUGGINGFACEHUB_API_TOKEN secret in your Space.")
     st.stop()
 
+# Import torch
+import torch
+
+# Set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 # Load the tokenizer and model with the token
 tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     use_auth_token=hf_token,
-    device_map='auto',
-    torch_dtype='auto'  # Adjust based on your environment
-)
+    device_map=None,  # We'll set the device manually
+    torch_dtype=torch.float32  # Use float32 to avoid half-precision issues
+).to(device)
 
 # Create the text-generation pipeline with appropriate parameters
 pipe = pipeline(
@@ -43,13 +50,13 @@ pipe = pipeline(
     repetition_penalty=1.1,
     do_sample=True,  # Use sampling to introduce some randomness
     eos_token_id=tokenizer.eos_token_id,
-    pad_token_id=tokenizer.eos_token_id
+    pad_token_id=tokenizer.eos_token_id,
+    device=0 if torch.cuda.is_available() else -1  # Use GPU if available
 )
 
 # Wrap the pipeline with HuggingFacePipeline for use in LangChain
 llm = HuggingFacePipeline(pipeline=pipe)
 
-# ... rest of your code ...
 
 
 # Step 1: Upload CSV data file (or use default)