Update app.py

app.py

@@ -6,7 +6,6 @@ import torch
 from transformers import AutoModelForCausalLM, GemmaTokenizerFast, pipeline
 from langchain_huggingface import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
 
 # Logging setup
 log_file = '/tmp/app_debug.log'
@@ -33,7 +32,7 @@ else:
     logger.warning("GPU is not available. Proceeding with CPU setup.")
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        device_map="cpu",
+        device_map="auto",
         low_cpu_mem_usage=True,
         token=os.getenv('HF_TOKEN'),
     )
@@ -92,9 +91,7 @@ def predict(message, chat_history=[]):
 # Gradio UI
 interface = gr.Interface(
     fn=predict,
-    inputs=[
-        gr.Textbox(label="User input")
-    ],
+    inputs=gr.Textbox(label="User input"),
     outputs="text",
     live=True,
 )