peft model logic added

app.py

@@ -1,50 +1,86 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel, PeftConfig
 import torch
+import os
+
+def load_model(model_id, model_type="base"):
+    try:
+        if model_type == "base":
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            return tokenizer, model
+        else:  # finetuned model with PEFT
+            # Load the base model first
+            base_model_id = "satyanayak/gemma-3-base"
+            tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+            base_model = AutoModelForCausalLM.from_pretrained(
+                base_model_id,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            
+            # Load and merge the PEFT adapters
+            model = PeftModel.from_pretrained(
+                base_model,
+                model_id,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            return tokenizer, model
+    except Exception as e:
+        print(f"Error loading {model_type} model: {str(e)}")
+        return None, None
 
 # Load base model and tokenizer
 base_model_id = "satyanayak/gemma-3-base"
-base_tokenizer = AutoTokenizer.from_pretrained(base_model_id)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+base_tokenizer, base_model = load_model(base_model_id, "base")
 
 # Load finetuned model and tokenizer
 finetuned_model_id = "satyanayak/gemma-3-GRPO"
-finetuned_tokenizer = AutoTokenizer.from_pretrained(finetuned_model_id)
-finetuned_model = AutoModelForCausalLM.from_pretrained(
-    finetuned_model_id,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+finetuned_tokenizer, finetuned_model = load_model(finetuned_model_id, "finetuned")
 
 def generate_base_response(prompt, max_length=512):
-    inputs = base_tokenizer(prompt, return_tensors="pt").to(base_model.device)
-    outputs = base_model.generate(
-        **inputs,
-        max_length=max_length,
-        num_return_sequences=1,
-        temperature=0.7,
-        do_sample=True,
-        pad_token_id=base_tokenizer.eos_token_id
-    )
-    response = base_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    if base_model is None or base_tokenizer is None:
+        return "Error: Base model failed to load. Please check if the model files are properly uploaded to Hugging Face."
+    
+    try:
+        inputs = base_tokenizer(prompt, return_tensors="pt").to(base_model.device)
+        outputs = base_model.generate(
+            **inputs,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=base_tokenizer.eos_token_id
+        )
+        response = base_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        return f"Error generating response with base model: {str(e)}"
 
 def generate_finetuned_response(prompt, max_length=512):
-    inputs = finetuned_tokenizer(prompt, return_tensors="pt").to(finetuned_model.device)
-    outputs = finetuned_model.generate(
-        **inputs,
-        max_length=max_length,
-        num_return_sequences=1,
-        temperature=0.7,
-        do_sample=True,
-        pad_token_id=finetuned_tokenizer.eos_token_id
-    )
-    response = finetuned_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    if finetuned_model is None or finetuned_tokenizer is None:
+        return "Error: Finetuned model failed to load. Please check if the model files are properly uploaded to Hugging Face."
+    
+    try:
+        inputs = finetuned_tokenizer(prompt, return_tensors="pt").to(finetuned_model.device)
+        outputs = finetuned_model.generate(
+            **inputs,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=finetuned_tokenizer.eos_token_id
+        )
+        response = finetuned_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        return f"Error generating response with finetuned model: {str(e)}"
 
 # Example prompts
 examples = [

requirements.txt

@@ -1,4 +1,5 @@
 gradio>=4.19.2
 transformers>=4.38.0
 torch>=2.2.0
-accelerate>=0.27.0 
+accelerate>=0.27.0
+peft>=0.9.0