Update app.py

app.py

@@ -1,97 +1,41 @@
-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)
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+import torch
 
 # 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 tokenizer(inputs, targets, padding="max_length", truncation=True)
+ds = load_dataset("knowrohit07/gita_dataset")
 
 # Load the model and tokenizer
-model_name = "t5-base"  # Or any other suitable model
+model_name = "deepset/roberta-base-squad2"
 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(
-    output_dir="./results",
-    evaluation_strategy="epoch",
-    save_strategy="epoch",
-    learning_rate=2e-5,
-    per_device_train_batch_size=16,
-    per_device_eval_batch_size=16,
-    num_train_epochs=3,
-    weight_decay=0.01,
-    logging_dir="./logs",  # Specify the logging directory
-)
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
 
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    data_collator=default_data_collator,
-    train_dataset=tokenized_dataset["train"],
-    eval_dataset=tokenized_dataset["validation"],
-)
-
-logging.info("Starting training...")
-trainer.train()
-logging.info("Training completed.")
-
-# Save the fine-tuned model
-model.save_pretrained("gita_model")
-tokenizer.save_pretrained("gita_tokenizer")
-
-# Define the question-answering function
 def answer_question(question):
-    """
-    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")
+    # Combine all text from the dataset
+    context = " ".join([item['Text'] for item in ds['train']])
 
-        # Preprocess the input
-        input_ids = tokenizer(question, return_tensors="pt").input_ids
+    # Tokenize input
+    inputs = tokenizer.encode_plus(question, context, return_tensors="pt", max_length=512, truncation=True)
 
-        # 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)
+    # Get model output
+    outputs = model(**inputs)
 
-        return answer.strip()
+    # Process the output to get the 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]))
 
-    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."
+    return answer
 
-# Create the Gradio interface
-interface = gr.Interface(
+# Define the Gradio interface
+iface = gr.Interface(
     fn=answer_question,
-    inputs="text",
+    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
     outputs="text",
     title="Bhagavad Gita Q&A",
-    description="Ask your questions about the Bhagavad Gita and receive insights from the model."
+    description="Ask a question about the Bhagavad Gita, and get an answer based on the dataset."
 )
 
-interface.launch()
+# Launch the app
+iface.launch()