dev

app.py

@@ -2,15 +2,15 @@ import spaces
 import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import DataCollatorForSeq2Seq, AutoConfig
-from datasets import load_dataset, concatenate_datasets, load_from_disk, DatasetDict
+#from datasets import load_dataset, concatenate_datasets, load_from_disk, DatasetDict
 import traceback
-from sklearn.metrics import accuracy_score
-import numpy as np
+#from sklearn.metrics import accuracy_score
+#import numpy as np
 import torch
 import os
-import evaluate
-from huggingface_hub import login
-from peft import get_peft_model, LoraConfig
+#import evaluate
+#from huggingface_hub import login
+#from peft import get_peft_model, LoraConfig
 
 os.environ['HF_HOME'] = '/data/.huggingface'
 '''
@@ -245,16 +245,19 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
 
 # Define Gradio interface
 @spaces.GPU
-def predict(text):
-    
+def predict(text):    
     config = AutoConfig.from_pretrained("shorecode/t5-efficient-tiny-nh8-summarizer")
     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)
-    outputs = model(inputs)
-    predictions = outputs.logits.argmax(dim=-1)
-    return predictions.item()
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+        inputs = {key: value.to('cuda') for key, value in inputs.items()}    
+    with torch.no_grad():  # Disable gradient calculation for inference
+        outputs = model.generate(inputs)
+    predictions = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return predictions
 
 
 @spaces.GPU(duration=120)