dev

app.py

@@ -243,7 +243,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             
         try:
             train_result = trainer.train(resume_from_checkpoint=True)
-        except:
+        except Exception as e:
+            print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
             import shutil
             checkpoint_dir = training_args.output_dir
             # If the trainer_state.json is missing, look for the previous checkpoint
@@ -266,6 +267,21 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
     return 'DONE!'#train_result
 
+@spaces.GPU
+def test(text):
+    from transformers import pipeline
+    model_name = 'shorecode/t5-efficient-tiny-nh8-summarizer'
+    summarizer = pipeline(
+        "summarization",
+        model=model_name,
+        tokenizer=model_name,
+        clean_up_tokenization_spaces=True, 
+    )        
+    
+    max_length = 500
+    summary = summarizer(text, max_length=max_length, min_length=40, no_repeat_ngram_size=2)
+    return summary 
+
 
 @spaces.GPU(duration=120)
 def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
@@ -288,6 +304,24 @@ def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     #model = get_peft_model(model, lora_config)
     result = fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad)    
     return result
+
+
+try:
+    iface = gr.Interface(
+        fn=test,
+        inputs=[
+            gr.Textbox(label="Text to summarize:"),
+        ],
+        outputs="text",
+        title="Fine-Tune Hugging Face Model shorecode/t5-efficient-tiny-nh8-summarizer",
+        description="This interface allows you to test shorecode/t5-efficient-tiny-nh8-summarizer."
+    )
+
+    # Launch the interface
+    iface.launch()  
+except Exception as e:
+    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
+'''
 # Create Gradio interface
 try:
     iface = gr.Interface(
@@ -309,4 +343,5 @@ try:
     # Launch the interface
     iface.launch()    
 except Exception as e:
-    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
+    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
+'''