updated app.py

app.py

@@ -6,7 +6,7 @@ import time
 
 # Configuration
 BASE_MODEL = "microsoft/phi-2"
-ADAPTER_MODEL = "pradeep6kumar2024/phi2-qlora-assistant"  # Your actual model ID
+ADAPTER_MODEL = "pradeep6kumar2024/phi2-qlora-assistant"
 
 class ModelWrapper:
     def __init__(self):
@@ -16,48 +16,76 @@ class ModelWrapper:
         
     def load_model(self):
         if not self.loaded:
-            print("Loading model and tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
-            base_model = AutoModelForCausalLM.from_pretrained(
-                BASE_MODEL,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                trust_remote_code=True
-            )
-            
-            print("Loading LoRA adapter...")
-            self.model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)
-            self.loaded = True
-            print("Model loading complete!")
+            try:
+                print("Loading tokenizer...")
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    BASE_MODEL,
+                    trust_remote_code=True,
+                    padding_side="left"
+                )
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                
+                print("Loading base model...")
+                base_model = AutoModelForCausalLM.from_pretrained(
+                    BASE_MODEL,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                    trust_remote_code=True,
+                    use_flash_attention_2=False  # Disable flash attention if causing issues
+                )
+                
+                print("Loading LoRA adapter...")
+                self.model = PeftModel.from_pretrained(
+                    base_model,
+                    ADAPTER_MODEL,
+                    torch_dtype=torch.float16,
+                    device_map="auto"
+                )
+                self.model.eval()
+                print("Model loading complete!")
+                self.loaded = True
+            except Exception as e:
+                print(f"Error during model loading: {str(e)}")
+                raise
     
-    def generate_response(self, prompt, max_length=512, temperature=0.7, top_p=0.9, stream=False):
+    def generate_response(self, prompt, max_length=512, temperature=0.7, top_p=0.9):
         if not self.loaded:
             self.load_model()
         
-        # Tokenize input
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-        
-        # Generate
-        start_time = time.time()
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_length=max_length,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=True,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id
-            )
-        
-        # Decode response
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        if response.startswith(prompt):
-            response = response[len(prompt):].strip()
-        
-        generation_time = time.time() - start_time
-        return response, generation_time
+        try:
+            # Tokenize input
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=512,
+                padding=True
+            ).to(self.model.device)
+            
+            # Generate
+            start_time = time.time()
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_length=max_length,
+                    temperature=temperature,
+                    top_p=top_p,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    repetition_penalty=1.1
+                )
+            
+            # Decode response
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            if response.startswith(prompt):
+                response = response[len(prompt):].strip()
+            
+            generation_time = time.time() - start_time
+            return response, generation_time
+        except Exception as e:
+            print(f"Error during generation: {str(e)}")
+            raise
 
 # Initialize model wrapper
 model_wrapper = ModelWrapper()
@@ -65,6 +93,9 @@ model_wrapper = ModelWrapper()
 def generate_text(prompt, max_length=512, temperature=0.7, top_p=0.9):
     """Gradio interface function"""
     try:
+        if not prompt.strip():
+            return "Please enter a prompt."
+        
         response, gen_time = model_wrapper.generate_response(
             prompt, 
             max_length=max_length,
@@ -73,7 +104,8 @@ def generate_text(prompt, max_length=512, temperature=0.7, top_p=0.9):
         )
         return f"Generated in {gen_time:.2f} seconds:\n\n{response}"
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        print(f"Error in generate_text: {str(e)}")
+        return f"Error generating response: {str(e)}\nPlease try again with a different prompt or parameters."
 
 # Create the Gradio interface
 demo = gr.Interface(
@@ -159,6 +191,5 @@ demo = gr.Interface(
     cache_examples=False
 )
 
-# Launch with sharing enabled
 if __name__ == "__main__":
     demo.launch()