Update app.py

app.py

@@ -2,28 +2,25 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 from huggingface_hub import InferenceClient
+import re
+import torch
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
 
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("sarvamai/sarvam-1")
-
-# Load the base model
-base_model = AutoModelForCausalLM.from_pretrained("sarvamai/sarvam-1")
-
-# Load the PEFT adapter
-peft_model = PeftModel.from_pretrained(base_model, "KGSAGAR/Sarvam-1-text-normalization-3r")
-peft_model = peft_model.merge_and_unload()
-
-
-# client = InferenceClient(peft_model)
-
-
-import re
-import torch
-from transformers import AutoTokenizer
+# Model and tokenizer loading (outside the respond function)
+try:
+    tokenizer = AutoTokenizer.from_pretrained("sarvamai/sarvam-1")
+    base_model = AutoModelForCausalLM.from_pretrained("sarvamai/sarvam-1")
+    peft_model = PeftModel.from_pretrained(base_model, "KGSAGAR/Sarvam-1-text-normalization-3r")
+    peft_model = peft_model.merge_and_unload()
+    print("Model loaded successfully!")  # Add this line
+except Exception as e:
+    print(f"Error loading model: {e}")
+    tokenizer = None
+    base_model = None
+    peft_model = None
 
 def respond(
     message,
@@ -32,8 +29,6 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    peft_model,
-    tokenizer_name,
 ):
     """
     Generates a response based on the user message and history using the provided PEFT model.
@@ -45,12 +40,14 @@ def respond(
         max_tokens (int): The maximum number of tokens to generate.
         temperature (float): The temperature parameter for generation.
         top_p (float): The top_p parameter for nucleus sampling.
-        peft_model: The pre-trained fine-tuned model for generation.
-        tokenizer_name (str): The name or path of the tokenizer.
 
     Yields:
         str: The generated response up to the current token.
     """
+    global tokenizer, peft_model #access global variables
+    if tokenizer is None or peft_model is None:
+        yield "Model loading failed. Please check the logs."
+        return
 
     # Construct the prompt
     prompt = system_message
@@ -63,14 +60,20 @@ def respond(
 
     # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+
     # Generate the output
-    outputs = peft_model.generate(
-        **inputs,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True  # Enable sampling for more diverse outputs
-    )
+    try:
+        outputs = peft_model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True  # Enable sampling for more diverse outputs
+        )
+    except Exception as e:
+         yield f"Generation error: {e}"
+         return
+
 
     # Decode the generated tokens
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)