Create src/agent.py

src/agent.py

@@ -0,0 +1,36 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+class CodingAgent:
+    def __init__(self, model_path):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = AutoModelForCausalLM.from_pretrained(model_path).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+    def generate_code(self, prompt, max_length=512, temperature=0.7, top_k=50, top_p=0.95):
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=max_length,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                do_sample=True,
+                num_return_sequences=1,
+            )
+        
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    def answer_coding_question(self, question):
+        prompt = f"As a coding assistant, please answer the following question:\n\nQuestion: {question}\n\nAnswer:"
+        return self.generate_code(prompt)
+
+    def explain_code(self, code):
+        prompt = f"Please explain the following code:\n\n```python\n{code}\n```\n\nExplanation:"
+        return self.generate_code(prompt)
+
+    def suggest_improvements(self, code):
+        prompt = f"Please suggest improvements for the following code:\n\n```python\n{code}\n```\n\nSuggestions:"
+        return self.generate_code(prompt)