app.py

app.py

@@ -1,15 +1,26 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # Load the smaller model and tokenizer
-model_name = "distilgpt2"
+model_name = "distilgpt2"  # A smaller model that should work with 16GB of RAM
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# Set the device to GPU if available, else use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=150, num_return_sequences=1)
+    # Encode the input prompt
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    
+    # Generate the output sequence
+    outputs = model.generate(inputs, max_length=150, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+    
+    # Decode the generated sequence
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
     return response
 
 # Set up Gradio interface
@@ -21,4 +32,5 @@ iface = gr.Interface(
     description="Enter your prompt related to Bitcoin or cryptocurrency."
 )
 
+# Launch the interface
 iface.launch()