Update app.py

app.py

@@ -222,6 +222,10 @@ from fastapi.middleware.cors import CORSMiddleware
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+COLLECTION_NAME = "arabic_rag_collection"
+QDRANT_URL = os.getenv("QDRANT_URL", "https://12efeef2-9f10-4402-9deb-f070977ddfc8.eu-central-1-0.aws.cloud.qdrant.io:6333")
+QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.Jb39rYQW2rSE9RdXrjdzKY6T1RF44XjdQzCvzFkjat4")
+
 # Load model and tokenizer
 # model_name = "FreedomIntelligence/Apollo-7B"
 # model_name = "emilyalsentzer/Bio_ClinicalBERT"
@@ -243,24 +247,43 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
-
 generation_config = GenerationConfig(
     max_new_tokens=150,
-    temperature=0.2,
+    temperature=0.7,
     top_k=20,
     do_sample=True,
-    top_p=0.7,
-    repetition_penalty=1.3,
+    top_p=0.8,
 )
 
 # Create generation pipeline
 pipe = TextGenerationPipeline(
     model=model,
     tokenizer=tokenizer,
-    device=model.device.index if torch.cuda.is_available() else "cpu"
+    device=model.device.index if torch.cuda.is_available() else -1
+)
+
+llm = HuggingFacePipeline(pipeline=pipe)
+
+embedding = HuggingFaceEmbeddings(model_name="Omartificial-Intelligence-Space/GATE-AraBert-v1")
+
+qdrant_client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
+
+vector_store = Qdrant(
+    client=qdrant_client,
+    collection_name=COLLECTION_NAME,
+    embeddings=embedding
 )
 
+retriever = vector_store.as_retriever(search_kwargs={"k": 3})
+
+# ----------------- RAG Chain ------------------
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    retriever=retriever,
+    chain_type="stuff"
+)
+
+
 # Prompt formatter based on language
 def generate_prompt(message):
     lang = detect(message)
@@ -277,17 +300,6 @@ def generate_prompt(message):
 Question: {message}
 Answer:"""
 
-# Chat function
-# @app.post("/ask")
-# def chat_fn(message):
-#     prompt = generate_prompt(message)
-#     response = pipe(prompt,
-#                     max_new_tokens=512,
-#                     temperature=0.7,
-#                     do_sample = True,
-#                     top_p=0.9)[0]['generated_text']
-#     answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
-#     return {"Answer": answer}
 
 executor = ThreadPoolExecutor()
 
@@ -320,28 +332,19 @@ async def chat_fn(query: Query):
         "Answer": answer
     }
 
-# @app.post("/ask")
-# async def chat_fn(query: Query):
-#     message = query.message
-#     logger.info(f"Received message: {message}")
-    
-#     prompt = generate_prompt(message)
-
-#     try:
-#         start_time = time.time()
-        
-#         loop = asyncio.get_event_loop()
-#         response = await loop.run_in_executor(
-#             executor,
-#             lambda: pipe(prompt, max_new_tokens=150, temperature=0.6, do_sample=True, top_p=0.8)[0]['generated_text']
-#         )
-        
-#         duration = time.time() - start_time
-#         logger.info(f"Model inference completed in {duration:.2f} seconds")
+@app.post("/ask-rag")
+async def chat_fn(query: Query):
+    message = query.message
+    logger.info(f"Received message: {message}")
 
-#         logger.info(f"Generated answer: {answer}")
-#         return {"Answer": answer}
+    # Run RAG inference in thread
+    loop = asyncio.get_event_loop()
+    response = await loop.run_in_executor(executor, lambda: qa_chain.run(message))
 
-#     except Exception as e:
-#         logger.error(f"Inference failed: {str(e)}")
-#         raise HTTPException(status_code=500, detail="Model inference TimeOut failed.")
+    answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
+    
+    
+    return {
+        "response": response,
+        "answer": answer
+    }