Update app.py

app.py

@@ -1,30 +1,42 @@
 import gradio as gr
-from inference_script import generate_outreach_message 
-import pandas as pd 
+from inference_script import generate_outreach_message, load_model_and_tokenizer
+import os
+
+HF_TOKEN = os.environ.get("HF_TOKEN") 
 
 try:
-    df = pd.read_csv('dataset.csv')
-except FileNotFoundError:
-    df = pd.DataFrame({}) 
+    model, tokenizer, device = load_model_and_tokenizer(
+        "KirubaLS/Outreach_model_6", 
+        "google/gemma-2b", 
+        hf_token=HF_TOKEN 
+    )
+except Exception as e:
+    print(f"Error loading model and tokenizer: {e}")
+    model = None
+    tokenizer = None
+    device = None
 
 
 def predict_message(industry, company, prospect, designation, client_info, my_info):
+    if model is None or tokenizer is None:
+        return "Error: Model not loaded. Please check the application logs."
+
     message = generate_outreach_message(industry, company, prospect, designation, client_info, my_info)
     return message
 
 iface = gr.Interface(
     fn=predict_message,
     inputs=[
-                gr.Textbox(label="Industry"),
-                gr.Textbox(label="Company"),
-                gr.Textbox(label="Prospect Name"),
-                gr.Textbox(label="Designation"),
-                gr.Textbox(label="Client Info"),
-                gr.Textbox(label="My Info"),
-            ],
-            outputs=gr.Textbox(label="Generated Outreach Message"),
-            title="Gemma Outreach Message Generator",
-            description="Generate personalized first-level outreach messages using a fine-tuned Gemma model."
-        )
+        gr.Textbox(label="Industry"),
+        gr.Textbox(label="Company"),
+        gr.Textbox(label="Prospect Name"),
+        gr.Textbox(label="Designation"),
+        gr.Textbox(label="Client Info"),
+        gr.Textbox(label="My Info"),
+    ],
+    outputs=gr.Textbox(label="Generated Outreach Message"),
+    title="Gemma Outreach Message Generator",
+    description="Generate personalized first-level outreach messages using a fine-tuned Gemma model."
+)
 
 iface.launch()