Update app.py

app.py

@@ -101,26 +101,53 @@ def generate_prompt(query: str, context_snippets: list) -> str:
     )
     return prompt
 
-def get_llm_response(prompt: str, model_name: str = "meta-llama/Llama-2-7b-chat-hf", max_new_tokens: int = None) -> str:
-    if max_new_tokens is None:
-        max_new_tokens = 1024 if is_detailed_query(prompt) else 256
+# def get_llm_response(prompt: str, model_name: str = "meta-llama/Llama-2-7b-chat-hf", max_new_tokens: int = None) -> str:
+#     if max_new_tokens is None:
+#         max_new_tokens = 1024 if is_detailed_query(prompt) else 256
 
-    torch.cuda.empty_cache()
+#     torch.cuda.empty_cache()
+
+#     if not os.path.exists("offload"):
+#         os.makedirs("offload")
+    
+#     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=HF_TOKEN)
+#     model = AutoModelForCausalLM.from_pretrained(
+#         model_name,
+#         device_map="auto",
+#         offload_folder="offload",  # Specify the folder where weights will be offloaded
+#         use_safetensors=False,
+#         trust_remote_code=True,
+#         torch_dtype=torch.float16,
+#         token=HF_TOKEN
+#     )
 
-    if not os.path.exists("offload"):
-        os.makedirs("offload")
     
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=HF_TOKEN)
+#     text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
+#     outputs = text_gen(prompt, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.7)
+#     full_response = outputs[0]['generated_text']
+    
+#     marker = "Answer:"
+#     if marker in full_response:
+#         answer = full_response.split(marker, 1)[1].strip()
+#     else:
+#         answer = full_response.strip()
+    
+#     return answer
+
+def get_llm_response(prompt: str, model_name: str = "EleutherAI/gpt-neo-125M", max_new_tokens: int = None) -> str:
+    if max_new_tokens is None:
+        max_new_tokens = 256  # You can adjust this value as needed.
+    
+    torch.cuda.empty_cache()
+    
+    # Load the tokenizer and model for GPT-Neo 125M.
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         device_map="auto",
-        offload_folder="offload",  # Specify the folder where weights will be offloaded
         use_safetensors=False,
-        trust_remote_code=True,
-        torch_dtype=torch.float16,
-        token=HF_TOKEN
+        torch_dtype=torch.float32  # Using default precision since model is small.
     )
-
     
     text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
     outputs = text_gen(prompt, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.7)
@@ -134,6 +161,7 @@ def get_llm_response(prompt: str, model_name: str = "meta-llama/Llama-2-7b-chat-
     
     return answer
 
+
 ############################################
 # Gradio Interface Functions
 ############################################