Update app.py

app.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi, hf_hub_download
 # Configuration constants
 MODEL_REPO = "Gajendra5490/Scrached_Trained_Model"
 CURRENT_USER = "gajendra82"
-CURRENT_UTC = "2025-05-06 15:44:35"
+CURRENT_UTC = "2025-05-06 15:52:11"
 
 def setup_logging():
     logging.basicConfig(
@@ -23,86 +23,70 @@ def setup_logging():
 
 logger = setup_logging()
 
-# Initialize Hugging Face API with token
-try:
-    token = os.environ.get('HF_TOKEN')
-    if not token:
-        logger.error("No HF_TOKEN found in environment variables")
-        raise ValueError("HF_TOKEN not found")
-    
-    api = HfApi(token=token)
-    logger.info("Successfully initialized Hugging Face API")
-except Exception as e:
-    logger.error(f"Failed to initialize Hugging Face API: {e}")
-    raise
-
-class EnhancedTokenizer:
-    def __init__(self, vocab):
-        self.vocab = vocab
-        self.special_tokens = {
-            "<user>": len(vocab),
-            "<assistant>": len(vocab) + 1,
-            "<sep>": len(vocab) + 2,
-            "<eos>": len(vocab) + 3
-        }
-        
-    def encode(self, text):
-        tokens = text.split()
-        return [self.vocab.get(token, 0) if token not in self.special_tokens 
-                else self.special_tokens[token] for token in tokens]
-        
-    def decode(self, ids):
-        reverse_vocab = {v: k for k, v in self.vocab.items()}
-        reverse_special = {v: k for k, v in self.special_tokens.items()}
-        return " ".join(reverse_vocab.get(id, reverse_special.get(id, "<unk>")) 
-                       for id in ids)
-
 class ImprovedTransformer(torch.nn.Module):
-    def __init__(self, vocab_size, d_model=512, nhead=8, 
-                 num_encoder_layers=6, num_decoder_layers=6,
-                 dim_feedforward=2048, dropout=0.1, max_seq_length=128):
+    def __init__(
+        self,
+        vocab_size,
+        d_model=512,
+        nhead=8,
+        num_encoder_layers=6,
+        num_decoder_layers=6,
+        dim_feedforward=2048,
+        dropout=0.1,
+        max_seq_length=128
+    ):
         super().__init__()
         
         self.d_model = d_model
         self.embedding = torch.nn.Embedding(vocab_size, d_model)
-        self.pos_encoder = torch.nn.Embedding(max_seq_length, d_model)
         
-        encoder_layer = torch.nn.TransformerEncoderLayer(
-            d_model=d_model,
-            nhead=nhead,
-            dim_feedforward=dim_feedforward,
-            dropout=dropout,
-            batch_first=True
-        )
-        decoder_layer = torch.nn.TransformerDecoderLayer(
+        # Position encoding
+        position = torch.arange(max_seq_length).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model))
+        pe = torch.zeros(max_seq_length, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+        
+        # Main transformer
+        self.transformer = torch.nn.Transformer(
             d_model=d_model,
             nhead=nhead,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
             dim_feedforward=dim_feedforward,
             dropout=dropout,
             batch_first=True
         )
         
-        self.transformer_encoder = torch.nn.TransformerEncoder(
-            encoder_layer, 
-            num_encoder_layers
-        )
-        self.transformer_decoder = torch.nn.TransformerDecoder(
-            decoder_layer, 
-            num_decoder_layers
-        )
-        
-        self.output = torch.nn.Linear(d_model, vocab_size)
+        # Output layers
+        self.output_layer = torch.nn.Linear(d_model, vocab_size)
+        self.norm = torch.nn.LayerNorm(d_model)
         
     def forward(self, src, tgt):
-        src_pos = torch.arange(0, src.size(1)).unsqueeze(0).to(src.device)
-        tgt_pos = torch.arange(0, tgt.size(1)).unsqueeze(0).to(tgt.device)
+        # Create masks
+        src_mask = self.transformer.generate_square_subsequent_mask(src.size(1)).to(src.device)
+        tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)
         
-        src = self.embedding(src) + self.pos_encoder(src_pos)
-        tgt = self.embedding(tgt) + self.pos_encoder(tgt_pos)
+        # Embeddings
+        src = self.embedding(src) * torch.sqrt(torch.tensor(self.d_model))
+        tgt = self.embedding(tgt) * torch.sqrt(torch.tensor(self.d_model))
+        
+        # Add positional encoding
+        src = src + self.pe[:src.size(1)].transpose(0, 1)
+        tgt = tgt + self.pe[:tgt.size(1)].transpose(0, 1)
+        
+        # Transform
+        output = self.transformer(
+            src,
+            tgt,
+            src_mask=src_mask,
+            tgt_mask=tgt_mask
+        )
         
-        memory = self.transformer_encoder(src)
-        output = self.transformer_decoder(tgt, memory)
-        return self.output(output)
+        # Output processing
+        output = self.norm(output)
+        return self.output_layer(output)
 
 class ModelInference:
     def __init__(self):
@@ -113,9 +97,12 @@ class ModelInference:
 
     def load_model(self):
         try:
-            # Download model and tokenizer using API
-            self.logger.info(f"Downloading from {MODEL_REPO}")
+            token = os.environ.get('HF_TOKEN')
+            if not token:
+                raise ValueError("HF_TOKEN not found in environment variables")
             
+            # Download files
+            self.logger.info(f"Downloading from {MODEL_REPO}")
             model_path = hf_hub_download(
                 repo_id=MODEL_REPO,
                 filename="model.pt",
@@ -128,75 +115,86 @@ class ModelInference:
                 token=token
             )
             
-            # Load model data
-            self.logger.info("Loading model...")
-            model_data = torch.load(
-                model_path,
-                map_location=self.device,
-                weights_only=False
-            )
-            
-            # Load tokenizer
+            # Load tokenizer first
             self.logger.info("Loading tokenizer...")
             with open(tokenizer_path, 'r', encoding='utf-8') as f:
                 tokenizer_data = json.load(f)
             
-            # Initialize model
-            model_config = model_data['model_config']
-            vocab_size = len(tokenizer_data['vocab']) + 4  # Add special tokens
+            # Initialize tokenizer
+            self.vocab = tokenizer_data['vocab']
+            self.special_tokens = {
+                "<user>": len(self.vocab),
+                "<assistant>": len(self.vocab) + 1,
+                "<sep>": len(self.vocab) + 2,
+                "<eos>": len(self.vocab) + 3
+            }
+            
+            # Load model
+            self.logger.info("Loading model...")
+            model_data = torch.load(
+                model_path,
+                map_location=self.device
+            )
             
+            # Initialize model with correct vocab size
+            vocab_size = len(self.vocab) + len(self.special_tokens)
             self.model = ImprovedTransformer(
                 vocab_size=vocab_size,
-                d_model=model_config.get('d_model', 512),
-                nhead=model_config.get('nhead', 8),
-                num_encoder_layers=model_config.get('num_encoder_layers', 6),
-                num_decoder_layers=model_config.get('num_decoder_layers', 6),
-                dim_feedforward=model_config.get('dim_feedforward', 2048),
-                dropout=model_config.get('dropout', 0.1),
-                max_seq_length=model_config.get('max_seq_length', 128)
+                d_model=512,
+                nhead=8,
+                num_encoder_layers=3,
+                num_decoder_layers=3,
+                dim_feedforward=2048
             ).to(self.device)
             
+            # Load state dict
             self.model.load_state_dict(model_data['model_state_dict'])
             self.model.eval()
             
-            # Initialize tokenizer
-            self.tokenizer = EnhancedTokenizer(tokenizer_data['vocab'])
-            
-            self.logger.info("Model and tokenizer loaded successfully")
+            self.logger.info("Model loaded successfully")
             
         except Exception as e:
             self.logger.error(f"Error loading model: {str(e)}")
             raise
 
+    def encode(self, text):
+        tokens = text.split()
+        return [self.vocab.get(token, 0) if token not in self.special_tokens 
+                else self.special_tokens[token] for token in tokens]
+
+    def decode(self, ids):
+        reverse_vocab = {v: k for k, v in self.vocab.items()}
+        reverse_special = {v: k for k, v in self.special_tokens.items()}
+        return " ".join(reverse_vocab.get(id, reverse_special.get(id, "<unk>")) 
+                       for id in ids)
+
     @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>")
+            # Tokenize
+            input_ids = self.encode(f"<user> {input_text} <sep>")
             input_tensor = torch.tensor([input_ids]).to(self.device)
             
             # Initialize response
-            response_ids = [self.tokenizer.special_tokens["<assistant>"]]
+            response_ids = [self.special_tokens["<assistant>"]]
+            response_tensor = torch.tensor([response_ids]).to(self.device)
             
-            # Generate response
-            for _ in range(150):  # max length
-                curr_output = self.model(
-                    input_tensor, 
-                    torch.tensor([response_ids]).to(self.device)
-                )
-                next_token = curr_output[0, -1].argmax().item()
+            # Generate
+            for _ in range(150):
+                output = self.model(input_tensor, response_tensor)
+                next_token = output[0, -1].argmax().item()
                 
-                if next_token == self.tokenizer.special_tokens["<eos>"]:
+                if next_token == self.special_tokens["<eos>"]:
                     break
                     
                 response_ids.append(next_token)
+                response_tensor = torch.tensor([response_ids]).to(self.device)
             
-            # Decode response
-            answer = self.tokenizer.decode(response_ids)
+            # Decode
+            answer = self.decode(response_ids)
             answer = answer.replace("<assistant>", "").replace("<eos>", "").strip()
             
             self.logger.info(f"Generated response: {answer}")
@@ -206,17 +204,10 @@ class ModelInference:
             self.logger.error(f"Error generating answer: {str(e)}")
             return f"Error generating answer: {str(e)}"
 
-# Initialize model globally
-try:
-    print("Initializing model...")
-    model = ModelInference()
-    print("Model initialized successfully")
-except Exception as e:
-    print(f"Error initializing model: {str(e)}")
-    model = None
+# Initialize model
+model = None
 
 def process_input(input_text):
-    """Process input through Gradio"""
     global model
     try:
         if model is None:
@@ -243,8 +234,6 @@ interface = gr.Interface(
     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(),
     examples=[
@@ -254,5 +243,5 @@ interface = gr.Interface(
     ]
 )
 
-# For Hugging Face Spaces
+# Launch
 interface.launch()