Upload assistant.py

assistant.py

@@ -0,0 +1,36 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+MODEL_NAME = "google/gemma-2b-it"
+
+class LegalEaseAssistant:
+    def __init__(self, model_name=MODEL_NAME):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            device_map="cpu",
+            load_in_8bit=True,
+            torch_dtype=torch.float16
+        )
+
+    def generate_response(self, text, task_type):
+        task_prompts = {
+            "simplify": f"Simplify the following legal text in clear, plain language:\n\n{text}\n\nSimplified explanation:",
+            "summary": f"Provide a concise summary of the following legal document:\n\n{text}\n\nSummary:",
+            "key_terms": f"Identify and explain key legal terms:\n\n{text}\n\nKey Terms:",
+            "risk": f"Perform a risk analysis:\n\n{text}\n\nRisk Assessment:"
+        }
+
+        prompt = task_prompts.get(task_type, f"Analyze the following text:\n\n{text}\n\nAnalysis:")
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        outputs = self.model.generate(
+            **inputs, 
+            max_new_tokens=300, 
+            num_return_sequences=1,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )
+
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response.split(prompt.split("\n\n")[-1])[-1].strip()