Update src/RAGSample.py

src/RAGSample.py

@@ -347,6 +347,29 @@ def setup_retriever(use_kaggle_data: bool = False, kaggle_dataset: Optional[str]
 #     # Create a chain combining the prompt template and LLM
 #     return prompt | llm | StrOutputParser()
 
+def initialize_biogpt():
+    try:
+        hf_pipeline = pipeline(
+            "text-generation",
+            model="microsoft/BioGPT",
+            tokenizer="microsoft/BioGPT",
+            max_new_tokens=150,
+            temperature=0.3,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            return_full_text=False,
+            truncation=True,
+            do_sample=True,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            pad_token_id=1,
+            eos_token_id=2,
+        )
+        print("BioGPT loaded successfully!")
+        return hf_pipeline
+    except Exception as e:
+        print(f"Error loading BioGPT: {e}")
+        return None
 
 def setup_rag_chain() -> Runnable:
     """Sets up the RAG chain with a prompt template and an LLM."""
@@ -369,18 +392,7 @@ Answer:
     )
     
     # Initialize a local Hugging Face model
-    hf_pipeline = pipeline(
-        "text-generation",
-        model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",  # Excellent for Q&A tasks
-        tokenizer="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
-        max_new_tokens=150,  # Generate only 150 new tokens instead of max_length
-        temperature=0.1,
-        device=0 if torch.cuda.is_available() else -1,
-        return_full_text=False,
-        truncation=True,  # Truncate input if too long
-        do_sample=True,   # Enable sampling for better responses
-        pad_token_id=50256  # Add padding token to avoid warnings
-    )
+    hf_pipeline = initialize_biogpt()
     
     # Wrap it in LangChain
     llm = HuggingFacePipeline(pipeline=hf_pipeline)