dev

app.py

@@ -250,9 +250,22 @@ def predict(text):
     model = AutoModelForSeq2SeqLM.from_config(config)
     #initialize_weights(model)
     tokenizer = AutoTokenizer.from_pretrained('shorecode/t5-efficient-tiny-nh8-summarizer')
-    inputs = tokenizer(text, padding='max_length', max_length=512, truncation=True)
+    inputs = tokenizer(text, return_tensors="pt", padding='max_length', max_length=512, truncation=True)
+
+    # Move model and inputs to GPU if available
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+        inputs = {key: value.to('cuda') for key, value in inputs.items()}
+
+    # Generate outputs
     with torch.no_grad():  # Disable gradient calculation for inference
         outputs = model.generate(inputs)
+
+    ## Decode the generated output
+    #predictions = tokenizer.decode(outputs[0], skip_special_tokens=True)    
+    #inputs = tokenizer(text, padding='max_length', max_length=512, truncation=True)
+    #with torch.no_grad():  # Disable gradient calculation for inference
+        #outputs = model.generate(inputs)
     predictions = tokenizer.decode(outputs[0], skip_special_tokens=True)
     print('xxxxxxxxxxxxxxxxxxxxxxx')
     print(predictions)