Update app.py

app.py

@@ -1,30 +1,47 @@
-from transformers import CodeLlamaForConditionalGeneration, CodeLlamaTokenizer
+from flask import Flask, request
+import transformers
+from transformers import pipeline
 
-# Load pre-trained model and tokenizer
-model_name = "codellama/CodeLlama-7b-hf"  # Assuming your CodeLlama model name
-tokenizer = CodeLlamaTokenizer.from_pretrained(model_name)
-model = CodeLlamaForConditionalGeneration.from_pretrained(model_name)
+app = Flask(__name__)
 
-# System message
-system_message = "You are a code teaching assistant named OmniCode created by Anusha K. Answer all the code related questions being asked."
+# Load Code-Llama model and tokenizer (outside function for efficiency)
+model_name = "codellama/CodeLlama-7b-hf"
+tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+model = pipeline("text-generation", model=model_name, temperature=1.0)  # Set temperature to 1
 
-def generate_response(prompt, max_length=150, temperature=1.0):
-    input_text = system_message + "\n" + prompt
-    input_ids = tokenizer.encode(input_text, return_tensors='pt')
+def answer_code_questions(question):
+  """
+  Answers code-related questions using Code-Llama as OmniCode.
 
-    # Generate response
-    output = model.generate(input_ids, 
-                            max_length=max_length, 
-                            temperature=temperature,
-                            pad_token_id=tokenizer.eos_token_id,
-                            num_return_sequences=1)
+  Args:
+      question (str): The code-related question to be answered.
 
-    # Decode and return the response
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
+  Returns:
+      str: The answer generated by Code-Llama.
+  """
+
+  # Create system message for OmniCode
+  system_message = f"<<SYS>>\nYou are a code teaching assistant named OmniCode created by Anusha K.\n<</SYS>>\n"
+
+  # Combine system message and user question
+  prompt = f"{system_message}\n{question}"
+
+  # Generate response using Code-Llama
+  try:
+    response = model(prompt, max_length=512, truncation=True)[0]["generated_text"]
+    return response.strip()
+  except Exception as e:
+    print(f"Error during generation: {e}")
+    return "I encountered an error while processing your question. Please try rephrasing it or providing more context."
+
+@app.route("/", methods=["POST"])
+def answer_question():
+  """
+  Handles user-submitted questions and returns answers via POST request.
+  """
+  question = request.form["question"]
+  answer = answer_code_questions(question)
+  return answer
 
 if __name__ == "__main__":
-    while True:
-        user_input = input("You: ")
-        response = generate_response(user_input)
-        print("OmniCode:", response)
+  app.run(debug=True)  # Set debug=False for production use