Create app.py

app.py

@@ -0,0 +1,32 @@
+from fastapi import FastAPI, Query
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+from retriever import retrieve_documents
+
+# Load Mistral 7B model
+MODEL_NAME = "mistralai/Mistral-7B-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", torch_dtype=torch.float16)
+
+# Create inference pipeline
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# FastAPI server
+app = FastAPI()
+
+@app.get("/")
+def read_root():
+    return {"message": "Mistral 7B RAG API is running!"}
+
+@app.get("/generate/")
+def generate_response(query: str = Query(..., title="User Query")):
+    # Retrieve relevant documents
+    retrieved_docs = retrieve_documents(query)
+
+    # Format prompt for RAG
+    prompt = f"Use the following information to answer:\n{retrieved_docs}\n\nUser: {query}\nAI:"
+    
+    # Generate response
+    output = generator(prompt, max_length=256, do_sample=True, temperature=0.7)[0]["generated_text"]
+    
+    return {"query": query, "response": output}