Update tinyllama_inference.py

tinyllama_inference.py

@@ -1,5 +1,4 @@
 import json
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 def load_model():
@@ -20,7 +19,7 @@ Solution: "{code}"
 Return ONLY valid JSON: {{"stars": number, "feedback": string}}
 Do not include any extra text outside the JSON.
 """
-    # Load the model and tokenizer.
+    # Load model and tokenizer (for simplicity, we load them per request; consider caching for performance)
     tokenizer, model = load_model()
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=150)
@@ -31,13 +30,13 @@ Do not include any extra text outside the JSON.
         result = {"stars": 0, "feedback": "Evaluation failed. Unable to parse AI response."}
     return result
 
-# For direct testing from the command line
+# For direct testing from the command line:
 if __name__ == "__main__":
     import sys
     if len(sys.argv) < 3:
         print(json.dumps({"error": "Please provide a question and code as arguments"}))
-        sys.exit(1)
-    question = sys.argv[1]
-    code = sys.argv[2]
-    result = evaluate_code(question, code)
-    print(json.dumps(result))
+    else:
+        question = sys.argv[1]
+        code = sys.argv[2]
+        result = evaluate_code(question, code)
+        print(json.dumps(result))