Update app.py

app.py

@@ -1,26 +1,30 @@
-from transformers import AutoTokenizer
-import transformers
-import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-model = "codellama/CodeLlama-7b-hf"
+# Load pre-trained model and tokenizer
+model_name = "codellama/CodeLlama-7b-hf"  
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
+# System message
+system_message = "You are a code teaching assistant named OmniCode created by Anusha K. Answer all the code related questions being asked."
 
-sequences = pipeline(
-    'import socket\n\ndef ping_exponential_backoff(host: str):',
-    do_sample=True,
-    top_k=10,
-    temperature=0.1,
-    top_p=0.95,
-    num_return_sequences=1,
-    eos_token_id=tokenizer.eos_token_id,
-    max_length=200,
-)
-for seq in sequences:
-    print(f"Result: {seq['generated_text']}")
+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')
+
+    # Generate response
+    output = model.generate(input_ids, 
+                            max_length=max_length, 
+                            temperature=temperature,
+                            pad_token_id=tokenizer.eos_token_id,
+                            num_return_sequences=1)
+
+    # Decode and return the response
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
+
+if __name__ == "__main__":
+    while True:
+        user_input = input("You: ")
+        response = generate_response(user_input)
+        print("OmniCode:", response)