Update app.py

app.py

@@ -1,36 +1,48 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import re
+import logging
+
+logging.basicConfig(level=logging.INFO)
 
 def initialize_pipeline(token):
     """Initialize the Hugging Face pipeline with authentication token."""
-    return pipeline("text-generation", model="riotu-lab/ArabianGPT-0.8B-IslamicQA", use_auth_token=token)
-
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("riotu-lab/ArabianGPT-0.8B-IslamicQA", use_auth_token=token)
+        model = AutoModelForCausalLM.from_pretrained("riotu-lab/ArabianGPT-0.8B-IslamicQA", use_auth_token=token)
+        return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    except Exception as e:
+        logging.error(f"Failed to load model with token. Error: {str(e)}")
+        raise
 
 def generate_response(token, question, penalty_alpha, do_sample, top_k, temperature, repetition_penalty, max_new_tokens):
     """Generate the response using the configured parameters and provided question."""
-    # Initialize the pipeline with the token
-    pipe = initialize_pipeline(token)
-    
-    # Apply the formatted prompt
-    formatted_prompt_text = f"user\n{question}\nassistant:"
-    
-    # Generate the response using the model pipeline with specified configuration
-    result = pipe(formatted_prompt_text, 
-                  num_return_sequences=1,
-                  do_sample=do_sample,
-                  top_k=top_k,
-                  temperature=temperature,
-                  repetition_penalty=repetition_penalty,
-                  max_new_tokens=max_new_tokens)
+    try:
+        # Initialize the pipeline with the token
+        pipe = initialize_pipeline(token)
+        
+        # Apply the formatted prompt
+        formatted_prompt_text = f"user\n{question}\nassistant:"
+        
+        # Generate the response using the model pipeline with specified configuration
+        result = pipe(formatted_prompt_text, 
+                      num_return_sequences=1,
+                      do_sample=do_sample,
+                      top_k=top_k,
+                      temperature=temperature,
+                      repetition_penalty=repetition_penalty,
+                      max_new_tokens=max_new_tokens)
 
-    # Assuming the result is in the correct format, extract the generated text
-    response_text = result[0]['generated_text'] if result else ""
-    
-    # Extract and clean the response
-    cleaned_response = extract_and_clean_response(response_text)
+        # Assuming the result is in the correct format, extract the generated text
+        response_text = result[0]['generated_text'] if result else ""
+        
+        # Extract and clean the response
+        cleaned_response = extract_and_clean_response(response_text)
 
-    return cleaned_response
+        return cleaned_response
+    except Exception as e:
+        logging.error(f"An error occurred during generation: {str(e)}")
+        return "An error occurred. Please check your input and try again."
 
 def extract_and_clean_response(text: str) -> str:
     """Extracts the part after 'assistant:' and cleans unwanted patterns from it."""
@@ -64,6 +76,5 @@ iface = gr.Interface(
     live=False  # Requires user to press "submit" to run
 )
 
-
 # Launch the app
 iface.launch()