Update inference.py

inference.py

@@ -1,57 +1,36 @@
+import os
+import faiss
 import torch
-from evo_model import EvoTransformer
-from transformers import AutoTokenizer
-
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-def load_model(model_path="evo_hellaswag.pt"):
-    model = EvoTransformer()
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    model.to(device)
-    model.eval()
-    return model
-
-evo_model = load_model()
-
-def get_evo_response(prompt, option1, option2):
-    inputs = [f"{prompt} {option1}", f"{prompt} {option2}"]
-    encoded = tokenizer(inputs, padding=True, truncation=True, return_tensors="pt").to(device)
-
-    with torch.no_grad():
-        logits = evo_model(encoded["input_ids"]).squeeze(-1)
-
-    probs = torch.softmax(logits, dim=0)
-    best = torch.argmax(probs).item()
-
-    explanations = [
-        f"🅰️ Option 1: {option1}\nConfidence: {probs[0]:.2f}",
-        f"🅱️ Option 2: {option2}\nConfidence: {probs[1]:.2f}"
-    ]
-
-    final = f"Evo suggests: Option {best + 1}\n\n{explanations[best]}"
-    return final
-
-def get_gpt_response(prompt, option1, option2):
-    import openai
-    import os
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-
-    full_prompt = (
-        f"Question: {prompt}\n"
-        f"Option 1: {option1}\n"
-        f"Option 2: {option2}\n"
-        "Which option makes more sense and why?"
-    )
-
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
-            messages=[
-                {"role": "user", "content": full_prompt}
-            ]
-        )
-        return response.choices[0].message["content"].strip()
-    except Exception as e:
-        return f"GPT Error: {e}"
+from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
+from PyPDF2 import PdfReader
+
+class RAGRetriever:
+    def __init__(self):
+        self.encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+        self.index = faiss.IndexFlatL2(384)
+        self.contexts = []
+        self.ids = []
+
+    def add_document(self, text):
+        sentences = text.split("\n")
+        clean_sentences = [s.strip() for s in sentences if s.strip()]
+        embeddings = self.encoder.encode(clean_sentences)
+        self.index.add(embeddings)
+        self.contexts.extend(clean_sentences)
+
+    def retrieve(self, query, top_k=3):
+        q_vec = self.encoder.encode([query])
+        D, I = self.index.search(q_vec, top_k)
+        return [self.contexts[i] for i in I[0]]
+
+def extract_text_from_file(file_path):
+    ext = os.path.splitext(file_path)[-1].lower()
+    if ext == ".txt":
+        with open(file_path, "r", encoding="utf-8") as f:
+            return f.read()
+    elif ext == ".pdf":
+        reader = PdfReader(file_path)
+        return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
+    else:
+        return ""