deve

app.py

@@ -242,19 +242,20 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
     except Exception as e:
         return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
     return 'DONE!'#train_result
-'''
+
 # Define Gradio interface
 def predict(text):
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained('shorecode/t5-efficient-tiny-nh8-summarizer', num_labels=2)
+    tokenizer = AutoTokenizer.from_pretrained('shorecode/t5-efficient-tiny-nh8-summarizer')
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
     outputs = model(inputs)
     predictions = outputs.logits.argmax(dim=-1)
     return predictions.item()
-'''
+
 
 @spaces.GPU(duration=120)
-def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+#def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+def run_train(text):
     def initialize_weights(model):
         for name, param in model.named_parameters():
             if 'encoder.block.0.layer.0.DenseReluDense.wi.weight' in name:  # Example layer
@@ -262,7 +263,7 @@ def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
             elif 'encoder.block.0.layer.0.DenseReluDense.wo.weight' in name:  # Another example layer
                 torch.nn.init.kaiming_normal_(param.data)  # Kaiming initialization
     
-    config = AutoConfig.from_pretrained("google/t5-efficient-tiny-nh8")
+    config = AutoConfig.from_pretrained("shorecode/t5-efficient-tiny-nh8-summarizer")
     model = AutoModelForSeq2SeqLM.from_config(config)
     initialize_weights(model)
     lora_config = LoraConfig(
@@ -272,10 +273,12 @@ def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
         bias="none"  # Bias handling
     )
     #model = get_peft_model(model, lora_config)
-    result = fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad)
+    #result = fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad)
+    result = predict(text)
     return result
 # Create Gradio interface
-try:    
+try:
+    '''
     iface = gr.Interface(
         fn=run_train,
         inputs=[
@@ -293,7 +296,7 @@ try:
     )
     '''
     iface = gr.Interface(
-        fn=predict,
+        fn=run_train,
         inputs=[
             gr.Textbox(label="Query"),
         ],
@@ -301,7 +304,7 @@ try:
         title="Fine-Tune Hugging Face Model",
         description="This interface allows you to test a fine-tune Hugging Face model."
     )
-    '''
+    
     # Launch the interface
     iface.launch()    
 except Exception as e: