Refactor app.py to integrate InferenceClient for response generation and update requirements.txt to include datasets

app.py

@@ -1,8 +1,9 @@
 import os
 import torch
-import numpy as np
 import faiss
-from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
+from huggingface_hub import InferenceClient
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 
@@ -15,7 +16,7 @@ embedding_model = AutoModel.from_pretrained(embedding_model_name)
 
 def embed_texts(texts):
     """Generate embeddings for a list of texts."""
-    inputs = embedding_tokenizer('query: ' + texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = embedding_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
         outputs = embedding_model(**inputs)
     # Use mean pooling to get embeddings
@@ -53,43 +54,35 @@ def load_documents(document_mapping, folder_path="Data"):
 documents = load_documents(document_mapping)
 print(f"Loaded {len(documents)} documents.")
 
-# Load your model and tokenizer
-generation_model_name = "Qwen/Qwen2.5-0.5B-Instruct"
-
-generation_tokenizer = AutoTokenizer.from_pretrained(generation_model_name)
-generation_model = AutoModelForCausalLM.from_pretrained(generation_model_name, trust_remote_code=True, device_map="cpu")
+secret = os.environ.getattribute("API_TOKEN")
+client = InferenceClient(api_key=secret)
 
 def generate_response(query, retrieved_docs):
-    """Generate a response using Flan-T5-Large based on retrieved documents."""
+    """Generate a response with streaming tokens using OpenVINO and TextIteratorStreamer."""
     context = " ".join(retrieved_docs)
-     # More natural prompt
     prompt = (
         f"<s>Répondez à la question suivante de manière concise en utilisant uniquement les informations pertinentes du contexte fourni.\n\n"
         f"Contexte : {context}\n\n"
         f"Question : {query}\n\n"
         f"Réponse :"
     )
+    
     messages = [
-    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
-    {"role": "user", "content": prompt}
+        {"role": "system", "content": "Vous êtes un modèle de langage avancé en français, conçu pour fournir des réponses claires, complètes, grammaticalement correctes, et utiles, tout en restant courtois."},
+	    {
+		"role": "user",
+		"content": prompt,
+    	}
     ]
-    text = generation_tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    model_inputs = generation_tokenizer([text], return_tensors="pt").to(generation_model.device)
 
-    generated_ids = generation_model.generate(
-        **model_inputs,
-        max_new_tokens=512
+    completion = client.chat.completions.create(
+        model="meta-llama/Llama-3.2-3B-Instruct", 
+    	messages=messages, 
+        max_tokens=500,
     )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
 
-    response = generation_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return response
+    return completion.choices[0].message.content
+
 # 6. Query and Retrieve Relevant Documents
 def retrieve_documents(query, k=3):
     """Retrieve the top-k most relevant documents."""
@@ -101,11 +94,10 @@ def rag_pipeline(query):
     """Complete RAG pipeline."""
     # Step 1: Retrieve relevant documents
     relevant_docs = retrieve_documents(query, 1)
-    print(f"Retrieved {len(relevant_docs)} relevant documents.")
-    print(relevant_docs)
     # Step 2: Generate a response using the retrieved documents
     response = generate_response(query, relevant_docs)
-
+    print("Query:", query)
+    print("Response:", response)
     return response
 
 app = FastAPI()

requirements.txt

@@ -6,5 +6,5 @@ uvicorn
 numpy
 faiss-cpu
 faiss-gpu
-numpy
+datasets
 accelerate