Update app.py

app.py

@@ -1,13 +1,28 @@
 # app.py
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import requests
 import gradio as gr
 import torch
 
-# Load the Hugging Face model and tokenizer in 8-bit precision
+# Check if CUDA is available
+if torch.cuda.is_available():
+    # Configure 8-bit quantization
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=True,
+        llm_int8_threshold=6.0
+    )
+else:
+    # Skip quantization if CUDA is not available
+    quantization_config = None
+
+# Load the Hugging Face model and tokenizer
 model_name = "gpt2"  # Smaller and faster model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, load_in_8bit=True, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    device_map="auto" if torch.cuda.is_available() else None
+)
 
 # Groq API configuration
 GROQ_API_KEY = "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"
@@ -32,7 +47,7 @@ def generate_smart_contract(language, requirements):
     prompt = f"Generate a {language} smart contract with the following requirements: {requirements}"
     
     # Use the Hugging Face model to generate code
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # Move inputs to GPU
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
     outputs = model.generate(**inputs, max_length=300)  # Reduced max_length
     generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)