Update utils.py

utils.py

@@ -1,23 +1,113 @@
-def prepare_context(self, df: pd.DataFrame) -> str:
-    """Prépare le contexte avec les données du DataFrame en s'assurant que toutes les valeurs sont traitées comme du texte."""
+import pandas as pd
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+class CSVAnalyzer:
+    def __init__(self):
+        self.model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+        try:
+            # Tokenizer initialization with specific configuration
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name,
+                trust_remote_code=True,
+                use_fast=False
+            )
+            
+            # Padding token configuration
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                self.tokenizer.padding_side = "right"
+
+            # Model initialization
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True
+            )
+            
+            # Ensure model knows the pad_token
+            self.model.config.pad_token_id = self.tokenizer.pad_token_id
+            
+        except Exception as e:
+            print(f"Initialisierungsfehler: {str(e)}")
+            raise
+
+    def prepare_context(self, df: pd.DataFrame) -> str:
+        """Bereitet den Kontext mit den DataFrame-Daten vor."""
+        try:
+            context = "Dateninhalt:\n\n"
+            
+            # Limit rows to avoid context overflow
+            max_rows = min(len(df), 50)  # Maximum 50 rows
+            
+            # Conversion sécurisée des données
+            for idx in range(max_rows):
+                row = df.iloc[idx]
+                row_text = ""
+                for col in df.columns:
+                    if pd.notna(row[col]):
+                        # Conversion explicite en chaîne de caractères
+                        value = str(row[col]).strip()
+                        row_text += f"{col}: {value}\n"
+                context += f"Eintrag {idx + 1}:\n{row_text}\n---\n"
+            
+            return context.strip()
+            
+        except Exception as e:
+            raise Exception(f"Fehler bei der Kontextvorbereitung: {str(e)}")
+
+    def generate_response(self, context: str, query: str) -> str:
+        """Generiert eine Antwort auf die Frage unter Verwendung des Kontexts."""
+        prompt = f"""<s>[INST] Du bist ein Assistent, der auf Datenanalyse spezialisiert ist in ein Facility Management Unternehmen.
+Datenkontext:
+{context}
+Frage: {query}
+Antworte präzise und knapp, basierend ausschließlich auf den bereitgestellten Informationen. Gib das betreffende E-Mail inklusive Datum und Absender an. Erstelle bei Bedarf Analyse-Tabellen, um die Informationen strukturiert darzustellen. [/INST]"""
+
+        try:
+            # Tokenization with explicit padding handling
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=2048,
+                pad_to_multiple_of=8,
+                return_attention_mask=True
+            ).to(self.model.device)
+            
+            # Response generation
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
+                    max_new_tokens=512,
+                    temperature=0.7,
+                    top_p=0.95,
+                    repetition_penalty=1.15,
+                    do_sample=True,
+                    num_beams=1,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
+                )
+            
+            # Response decoding and cleaning
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            response = response.split("[/INST]")[-1].strip()
+            
+            return response
+            
+        except Exception as e:
+            return f"Generierungsfehler: {str(e)}"
+
+def analyze_csv(df: pd.DataFrame, query: str) -> str:
+    """Hauptfunktion zur CSV-Analyse und Fragenbeantwortung."""
     try:
-        context = "Dateninhalt:\n\n"
-        
-        # Limiter le nombre de lignes pour éviter un dépassement de contexte
-        max_rows = min(len(df), 50)  # Maximum 50 lignes
-        
-        # Conversion sécurisée des données
-        for idx in range(max_rows):
-            row = df.iloc[idx]
-            row_text = ""
-            for col in df.columns:
-                if pd.notna(row[col]):
-                    # Conversion explicite en chaîne de caractères
-                    value = str(row[col]).strip()
-                    row_text += f"{col}: {value}\n"
-            context += f"Eintrag {idx + 1}:\n{row_text}\n---\n"
-            
-        return context.strip()
+        analyzer = CSVAnalyzer()
+        context = analyzer.prepare_context(df)
+        response = analyzer.generate_response(context, query)
+        return response
         
     except Exception as e:
-        raise Exception(f"Fehler bei der Kontextvorbereitung: {str(e)}")
+        return f"Fehler bei der Analyse: {str(e)}"