Update app.py

app.py

@@ -6,48 +6,48 @@ def greet(name):
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()
 '''
+import os
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModel
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-
+from transformers import AutoTokenizer, AutoModel, T5ForConditionalGeneration, T5Tokenizer, GPT2LMHeadModel, GPT2Tokenizer
 import faiss
 import numpy as np
-import os
-
+import pytesseract
+from PIL import Image
+import gradio as gr
 
 # Define the path to the docs folder
 docs_path = "docs"
 
-# Load documents from the folder
+# Load images from the folder and extract text using pytesseract
 def load_documents(docs_path):
-    data_files = {"train": [os.path.join(docs_path, file) for file in os.listdir(docs_path) if file.endswith(".txt")]}
-    dataset = load_dataset("text", data_files=data_files)
-    return dataset["train"]
-
-documents_dataset = load_documents(docs_path)
+    documents = []
+    for filename in os.listdir(docs_path):
+        if filename.endswith((".png", ".jpg", ".jpeg")):
+            image_path = os.path.join(docs_path, filename)
+            text = pytesseract.image_to_string(Image.open(image_path))
+            documents.append(text)
+    return documents
 
+documents = load_documents(docs_path)
 
-
-
-# Load model and tokenizer
+# Load model and tokenizer for encoding documents
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 model = AutoModel.from_pretrained("distilbert-base-uncased")
 
 # Preprocess and encode documents
 def encode_documents(documents):
-    inputs = tokenizer(documents["text"], return_tensors='pt', padding=True, truncation=True)
+    inputs = tokenizer(documents, return_tensors='pt', padding=True, truncation=True)
     outputs = model(**inputs)
     embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()
     return embeddings
 
-embeddings = encode_documents(documents_dataset)
+embeddings = encode_documents(documents)
 
-# Index embeddings
+# Index embeddings using FAISS
 index = faiss.IndexFlatL2(embeddings.shape[1])
 index.add(embeddings)
 
-
-# Load T5 model and tokenizer
+# Load T5 model and tokenizer for question generation
 t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
 t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")
 
@@ -59,14 +59,11 @@ def generate_questions(text):
     return question
 
 def retrieve_documents(query, index, documents):
-    query_embedding = encode_documents({"text": [query]})
+    query_embedding = encode_documents([query])
     D, I = index.search(query_embedding, k=5)
     return [documents[i] for i in I[0]]
 
-# Answer Generation
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-
-# Load GPT-2 model and tokenizer
+# Load GPT-2 model and tokenizer for answer generation
 gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2")
 
@@ -77,56 +74,19 @@ def generate_answer(question, context):
     answer = gpt2_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return answer
 
-#fine tuning the language model
-
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
-
-# Load dataset
-dataset = load_dataset("text", data_files={"train": "docs/*.txt"})
-
-# Load model and tokenizer
-model_name = "gpt2"
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-model = GPT2LMHeadModel.from_pretrained(model_name)
-
-# Tokenize the dataset
-def tokenize_function(examples):
-    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
-
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
-
-# Set training arguments
-training_args = TrainingArguments(
-    output_dir="./results",
-    evaluation_strategy="epoch",
-    learning_rate=2e-5,
-    per_device_train_batch_size=2,
-    num_train_epochs=3,
-    weight_decay=0.01,
-)
-
-# Initialize Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets["train"],
-)
-
-# Fine-tune the model
-trainer.train()
+# Fine-tuning the language model (example code provided earlier)
 
+# Generate documents based on prompts
 def generate_document(prompt):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    outputs = model.generate(input_ids, max_length=512, num_return_sequences=1)
-    document = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    input_ids = gpt2_tokenizer.encode(prompt, return_tensors="pt")
+    outputs = gpt2_model.generate(input_ids, max_length=512, num_return_sequences=1)
+    document = gpt2_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return document
 
-# Gradio
-import gradio as gr
-
+# Gradio Interface
 def rag_pipeline(query):
     question = generate_questions(query)
-    retrieved_docs = retrieve_documents(question, index, documents_dataset["text"])
+    retrieved_docs = retrieve_documents(question, index, documents)
     context = " ".join(retrieved_docs)
     answer = generate_answer(question, context)
     return question, context, answer