Update app.py

app.py

@@ -35,6 +35,12 @@ except Exception as e:
     logger.error(f"Error loading model or tokenizer: {str(e)}")
     sys.exit(1)
 
+def clean_answer(answer):
+    # Remove special tokens and leading/trailing whitespace
+    special_tokens = set(tokenizer.all_special_tokens)
+    cleaned_answer = ' '.join(token for token in answer.split() if token not in special_tokens)
+    return cleaned_answer.strip()
+
 def answer_question(question, system_prompt, temperature, max_new_tokens, top_p, frequency_penalty, presence_penalty, top_k, echo, best_of):
     logger.info(f"Received question: {question}")
     logger.info(f"Parameters: temp={temperature}, max_tokens={max_new_tokens}, top_p={top_p}, freq_penalty={frequency_penalty}, pres_penalty={presence_penalty}, top_k={top_k}, echo={echo}, best_of={best_of}")
@@ -55,9 +61,14 @@ def answer_question(question, system_prompt, temperature, max_new_tokens, top_p,
         logger.info("Processing output to get answer")
         answer_start = torch.argmax(outputs.start_logits)
         answer_end = torch.argmax(outputs.end_logits) + 1
-        answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
+        raw_answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
+        answer = clean_answer(raw_answer)
         logger.info(f"Generated answer: {answer}")
 
+        if not answer:
+            logger.warning("Generated answer was empty after cleaning")
+            answer = "I'm sorry, but I couldn't find a specific answer to that question based on the Bhagavad Gita. Could you please rephrase your question or ask about one of the core concepts like dharma, karma, bhakti, or the different types of yoga discussed in the Gita?"
+
         disclaimer = "\n\nPlease note: This response is generated by an AI model based on the Bhagavad Gita. For authoritative information, please consult the original text or scholarly sources."
         full_response = answer + disclaimer
         logger.info("Answer generated successfully")