Update app.py

app.py

@@ -1,3 +1,4 @@
+# server.py
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
@@ -30,6 +31,14 @@ app.add_middleware(
 BASE_MODEL_PATH = "HuggingFaceTB/SmolLM2-135M-Instruct"
 ADAPTER_PATH = "khurrameycon/SmolLM-135M-Instruct-qa_pairs_converted.json-25epochs"
 
+def format_prompt(instruction):
+    """Format the prompt according to the model's expected format."""
+    return f"""### Instruction:
+{instruction}
+
+### Response:
+"""
+
 def load_model_and_tokenizer():
     """Load the model, tokenizer, and adapter weights."""
     try:
@@ -38,11 +47,28 @@ def load_model_and_tokenizer():
             BASE_MODEL_PATH,
             torch_dtype=torch.float16,
             trust_remote_code=True,
-            device_map="auto"
+            device_map="auto",
+            use_cache=True
         )
 
         logger.info("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(
+            BASE_MODEL_PATH,
+            padding_side="left",
+            truncation_side="left"
+        )
+        
+        # Ensure the tokenizer has the necessary special tokens
+        special_tokens = {
+            "pad_token": "<|padding|>",
+            "eos_token": "</s>",
+            "bos_token": "<s>",
+            "unk_token": "<|unknown|>"
+        }
+        tokenizer.add_special_tokens(special_tokens)
+        
+        # Resize the model embeddings to match the new tokenizer size
+        model.resize_token_embeddings(len(tokenizer))
 
         logger.info("Downloading adapter weights...")
         adapter_path_local = snapshot_download(repo_id=ADAPTER_PATH)
@@ -71,41 +97,58 @@ except Exception as e:
 def generate_response(model, tokenizer, instruction, max_new_tokens=2048):
     """Generate a response from the model based on an instruction."""
     try:
-        logger.info(f"Generating response for instruction: {instruction[:100]}...")
+        # Format the prompt
+        formatted_prompt = format_prompt(instruction)
+        logger.info(f"Formatted prompt: {formatted_prompt}")
         
         # Encode input with truncation
-        inputs = tokenizer.encode(
-            instruction,
+        inputs = tokenizer(
+            formatted_prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=tokenizer.model_max_length
+            max_length=tokenizer.model_max_length,
+            padding=True,
+            add_special_tokens=True
         ).to(model.device)
 
-        logger.info(f"Input shape: {inputs.shape}")
+        logger.info(f"Input shape: {inputs.input_ids.shape}")
         
-        # Create attention mask
-        attention_mask = torch.ones(inputs.shape, device=model.device)
-
         # Generate response
-        outputs = model.generate(
-            inputs,
-            attention_mask=attention_mask,
-            max_new_tokens=max_new_tokens,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-        )
+        with torch.inference_mode():
+            outputs = model.generate(
+                input_ids=inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                max_new_tokens=max_new_tokens,
+                temperature=0.7,
+                top_p=0.9,
+                top_k=50,
+                do_sample=True,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1,
+                length_penalty=1.0,
+                no_repeat_ngram_size=3
+            )
 
         logger.info(f"Output shape: {outputs.shape}")
 
-        # Decode and strip input prompt from response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        generated_text = response[len(instruction):].strip()
+        # Decode the response
+        response = tokenizer.decode(
+            outputs[0, inputs.input_ids.shape[1]:],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True
+        )
 
-        logger.info(f"Generated text length: {len(generated_text)}")
-        return generated_text
+        response = response.strip()
+        logger.info(f"Generated text length: {len(response)}")
+        logger.info(f"Generated text preview: {response[:100]}...")
+        
+        if not response:
+            logger.warning("Empty response generated")
+            raise ValueError("Model generated an empty response")
+            
+        return response
     except Exception as e:
         logger.error(f"Error generating response: {e}", exc_info=True)
         raise ValueError(f"Error generating response: {e}")
@@ -127,11 +170,6 @@ async def generate_text(input: ModelInput, request: Request):
             max_new_tokens=input.max_new_tokens
         )
         
-        if not response:
-            logger.warning("Generated empty response")
-            return {"generated_text": "", "warning": "Empty response generated"}
-            
-        logger.info(f"Generated response length: {len(response)}")
         return {"generated_text": response}
     except Exception as e:
         logger.error(f"Error in generate_text endpoint: {e}", exc_info=True)
@@ -148,5 +186,10 @@ async def health_check():
     return {
         "status": "healthy",
         "model_loaded": model is not None and tokenizer is not None,
-        "model_device": str(next(model.parameters()).device) if model else None
-    }
+        "model_device": str(next(model.parameters()).device) if model else None,
+        "tokenizer_vocab_size": len(tokenizer) if tokenizer else None
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")