Update app.py

app.py

@@ -1,25 +1,21 @@
-try:
-    import torch
-    import gradio as gr
-    import json
-    import logging
-    from pathlib import Path
-except ImportError as e:
-    print(f"Missing required package: {e}")
-    print("Please install required packages using:")
-    print("pip install torch gradio")
-    exit(1)
+import torch
+import gradio as gr
+import json
+import logging
+from pathlib import Path
+from huggingface_hub import hf_hub_download
 
 # Configuration constants
+MODEL_REPO = "Gajendra5490/Scrached_Trained_Model"
 CURRENT_USER = "gajendra82"
-CURRENT_UTC = "2025-05-06 15:10:08"
+CURRENT_UTC = "2025-05-06 15:15:08"
 
 def setup_logging():
     logging.basicConfig(
         level=logging.INFO,
         format='%(asctime)s - %(levelname)s - %(message)s',
         handlers=[
-            logging.StreamHandler()  # Only console logging for Hugging Face
+            logging.StreamHandler()
         ]
     )
     return logging.getLogger(__name__)
@@ -30,15 +26,26 @@ class ModelInference:
     def __init__(self):
         self.logger = logging.getLogger(__name__)
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.logger.info(f"Using device: {self.device}")
         self.load_model()
-        
+
     def load_model(self):
         try:
-            # Load model - using local path in Hugging Face space
-            model_path = Path('model.pt')
-            tokenizer_path = Path('tokenizer.json')
+            # Download model and tokenizer from your Hugging Face repository
+            self.logger.info(f"Downloading model from {MODEL_REPO}")
+            
+            model_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename="model.pt"
+            )
             
-            # Load model with weights_only=False
+            tokenizer_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename="tokenizer.json"
+            )
+            
+            # Load model
+            self.logger.info("Loading model...")
             model_data = torch.load(
                 model_path,
                 map_location=self.device,
@@ -46,6 +53,7 @@ class ModelInference:
             )
             
             # Load tokenizer
+            self.logger.info("Loading tokenizer...")
             with open(tokenizer_path, 'r', encoding='utf-8') as f:
                 tokenizer_data = json.load(f)
             
@@ -71,7 +79,7 @@ class ModelInference:
             from tokenizer import EnhancedTokenizer
             self.tokenizer = EnhancedTokenizer(tokenizer_data['vocab'])
             
-            self.logger.info("Model loaded successfully")
+            self.logger.info("Model and tokenizer loaded successfully")
             
         except Exception as e:
             self.logger.error(f"Error loading model: {e}")
@@ -80,14 +88,18 @@ class ModelInference:
     @torch.no_grad()
     def generate_answer(self, input_text: str) -> str:
         try:
+            # Clean input
+            input_text = input_text.strip()
+            self.logger.info(f"Processing input: {input_text}")
+            
             # Tokenize input
             input_ids = self.tokenizer.encode(f"<user> {input_text} <sep>")
             input_tensor = torch.tensor([input_ids]).to(self.device)
             
-            # Initialize response with start token
+            # Initialize response
             response_ids = [self.tokenizer.special_tokens["<assistant>"]]
             
-            # Generate output
+            # Generate response
             for _ in range(150):  # max length
                 curr_output = self.model(
                     input_tensor, 
@@ -100,41 +112,32 @@ class ModelInference:
                     
                 response_ids.append(next_token)
             
-            # Decode output
+            # Decode response
             answer = self.tokenizer.decode(response_ids)
             answer = answer.replace("<assistant>", "").replace("<eos>", "").strip()
             
+            self.logger.info(f"Generated response: {answer}")
             return answer
             
         except Exception as e:
             self.logger.error(f"Error generating answer: {e}")
-            return f"Error generating answer: {str(e)}"
+            return f"Error: {str(e)}"
 
-# Initialize model globally for Hugging Face
+# Initialize model
 try:
-    print("Loading model...")
+    print("Initializing model...")
     model = ModelInference()
-    print("Model loaded successfully")
+    print("Model initialized successfully")
 except Exception as e:
-    print(f"Error loading model: {e}")
+    print(f"Error initializing model: {e}")
     model = None
 
 def process_input(input_text):
     """Process input through Gradio"""
     try:
         if model is None:
-            return "Error: Model not loaded properly"
-        
-        # Log the input
-        logger.info(f"Input received: {input_text}")
-        
-        # Generate answer
-        answer = model.generate_answer(input_text)
-        
-        # Log the output
-        logger.info(f"Generated answer: {answer}")
-        
-        return answer
+            return "Error: Model not initialized properly"
+        return model.generate_answer(input_text)
     except Exception as e:
         logger.error(f"Error processing input: {e}")
         return f"Error: {str(e)}"
@@ -143,21 +146,29 @@ def process_input(input_text):
 interface = gr.Interface(
     fn=process_input,
     inputs=gr.Textbox(
-        label="Input",
-        placeholder="Enter your input here...",
+        label="Input Question",
+        placeholder="Enter your question here...",
         lines=2
     ),
     outputs=gr.Textbox(
-        label="Answer",
+        label="Model Response",
         lines=4
     ),
     title="Model Inference Interface",
     description=f"""
+    Model Repository: {MODEL_REPO}
     Current User: {CURRENT_USER}
     Last Updated: {CURRENT_UTC} UTC
+    
+    Enter your question and click submit to get a response.
     """,
     theme=gr.themes.Soft(),
-    allow_flagging="never"
+    allow_flagging="never",
+    examples=[
+        ["What is this about?"],
+        ["Can you explain the topic?"],
+        ["Give me more details."]
+    ]
 )
 
 # For Hugging Face Spaces