Update app.py

app.py

@@ -1,47 +1,30 @@
-from flask import Flask, request
-import transformers
-from transformers import pipeline
+from transformers import LlamaForCausalLM, CodeLlamaTokenizer
 
-app = Flask(__name__)
+# Load pre-trained model and tokenizer
+model_name = "codellama/CodeLlama-7b-hf"  # You can replace it with your model name
+tokenizer = CodeLlamaTokenizer.from_pretrained(model_name)
+model = LlamaForCausalLM.from_pretrained(model_name)
 
-# 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
+# System message
+system_message = "You are a code teaching assistant named OmniCode created by Anusha K. Answer all the code related questions being asked."
 
-def answer_code_questions(question):
-  """
-  Answers code-related questions using Code-Llama as OmniCode.
+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')
 
-  Args:
-      question (str): The code-related question to be answered.
+    # Generate response
+    output = model.generate(input_ids, 
+                            max_length=max_length, 
+                            temperature=temperature,
+                            pad_token_id=tokenizer.eos_token_id,
+                            num_return_sequences=1)
 
-  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
+    # Decode and return the response
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
 
 if __name__ == "__main__":
-  app.run(debug=True)  # Set debug=False for production use
+    while True:
+        user_input = input("You: ")
+        response = generate_response(user_input)
+        print("OmniCode:", response)