Update app.py

app.py

@@ -1,23 +1,30 @@
+import logging
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TrainingArguments, Trainer
 from datasets import load_dataset
+import gradio as gr
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
 
 # Load the dataset
 dataset = load_dataset("knowrohit07/gita_dataset")
+logging.info("Dataset loaded successfully.")
 
 # Preprocess the dataset
 def preprocess_function(examples):
     inputs = [f"Question: {q} Answer:" for q in examples["question"]]
     targets = examples["answer"]
-    return {"input_ids": tokenizer(inputs, padding="max_length", truncation=True)["input_ids"],
-            "labels": tokenizer(targets, padding="max_length", truncation=True)["input_ids"]}
+    return tokenizer(inputs, targets, padding="max_length", truncation=True)
 
 # Load the model and tokenizer
 model_name = "t5-base"  # Or any other suitable model
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+logging.info("Model and tokenizer loaded successfully.")
 
 # Tokenize the dataset
 tokenized_dataset = dataset.map(preprocess_function, batched=True)
+logging.info("Dataset tokenized successfully.")
 
 # Fine-tune the model on the dataset
 training_args = TrainingArguments(
@@ -29,6 +36,7 @@ training_args = TrainingArguments(
     per_device_eval_batch_size=16,
     num_train_epochs=3,
     weight_decay=0.01,
+    logging_dir="./logs",  # Specify the logging directory
 )
 
 trainer = Trainer(
@@ -39,30 +47,51 @@ trainer = Trainer(
     eval_dataset=tokenized_dataset["validation"],
 )
 
+logging.info("Starting training...")
 trainer.train()
+logging.info("Training completed.")
 
-# Define the Gradio interface
+# Save the fine-tuned model
+model.save_pretrained("gita_model")
+tokenizer.save_pretrained("gita_tokenizer")
+
+# Define the question-answering function
 def answer_question(question):
-  """
-  This function takes a question about the Bhagavad Gita and uses the fine-tuned model to generate an answer.
+    """
+    Answers a question about the Bhagavad Gita using a fine-tuned model.
+
+    Args:
+        question: The question to be answered.
+
+    Returns:
+        The answer generated by the model.
+    """
+
+    try:
+        # Load the fine-tuned model and tokenizer
+        model = AutoModelForSeq2SeqLM.from_pretrained("gita_model")
+        tokenizer = AutoTokenizer.from_pretrained("gita_tokenizer")
+
+        # Preprocess the input
+        input_ids = tokenizer(question, return_tensors="pt").input_ids
+
+        # Generate the answer
+        output = model.generate(input_ids, max_length=500, no_repeat_ngram_size=2)
+        answer = tokenizer.decode(output[0], skip_special_tokens=True)
 
-  Args:
-      question: A string representing the user's question.
+        return answer.strip()
 
-  Returns:
-      A string representing the model's generated answer.
-  """
-  input_ids = tokenizer(question, return_tensors="pt").input_ids
-  output = model.generate(input_ids, max_length=500, no_repeat_ngram_size=2)
-  answer = tokenizer.decode(output[0], skip_special_tokens=True)
-  return answer.strip()
+    except Exception as e:
+        logging.error(f"An error occurred: {e}")
+        return "I couldn't find an answer to your question. Please try rephrasing it or asking something different."
 
+# Create the Gradio interface
 interface = gr.Interface(
-  fn=answer_question,
-  inputs="text",
-  outputs="text",
-  title="Bhagavad Gita Q&A",
-  description="Ask your questions about the Bhagavad Gita and receive insights from the model."
+    fn=answer_question,
+    inputs="text",
+    outputs="text",
+    title="Bhagavad Gita Q&A",
+    description="Ask your questions about the Bhagavad Gita and receive insights from the model."
 )
 
 interface.launch()