Update inference.py

inference.py

@@ -1,36 +1,78 @@
-import os
-import faiss
 import torch
-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 ""
+from evo_model import EvoTransformer
+from transformers import AutoTokenizer, pipeline
+from rag_utils import RAGRetriever, extract_text_from_file
+import os
+
+# Load Evo model
+def load_evo_model(model_path="evo_hellaswag.pt", device=None):
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model = EvoTransformer()
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
+    model.eval()
+    return model, device
+
+evo_model, device = load_evo_model()
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+
+# Load GPT-3.5 (via OpenAI API)
+import openai
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
+# RAG Retriever
+retriever = RAGRetriever()
+
+def get_context_from_file(file_obj):
+    file_path = file_obj.name
+    text = extract_text_from_file(file_path)
+    retriever.add_document(text)
+    return text
+
+# Evo prediction
+def get_evo_response(prompt, file=None):
+    # Step 1: augment context if document is uploaded
+    context = ""
+    if file is not None:
+        context_list = retriever.retrieve(prompt)
+        context = "\n".join(context_list)
+
+    full_prompt = f"{prompt}\n{context}"
+
+    # Step 2: use Evo to predict
+    options = ["Yes, proceed with the action.", "No, maintain current strategy."]
+    inputs = [f"{full_prompt} {opt}" for opt in options]
+
+    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()
+
+    return f"Evo suggests: {options[best]} (Confidence: {probs[best]:.2f})"
+
+# GPT-3.5 response
+def get_gpt_response(prompt, file=None):
+    context = ""
+    if file is not None:
+        context_list = retriever.retrieve(prompt)
+        context = "\n".join(context_list)
+
+    full_prompt = (
+        f"Question: {prompt}\n"
+        f"Relevant Context:\n{context}\n"
+        f"Answer like a financial advisor."
+    )
+
+    response = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {"role": "user", "content": full_prompt}
+        ],
+        temperature=0.4,
+    )
+
+    return response.choices[0].message.content.strip()